scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

refactor(auth): split key authentication from authorization

Browse source 
keys() now answers only who the credential belongs to, returning an identity or
the reason the key is unusable, and the new authorize() facet decides whether
that key may make the request: master-key-only routes, model access, team
member, key and per-model budgets, and the proxy budget alert. outcome.py is
split into identity.py (identities and authentication denials), outcome.py
(the authorization verdict) and validators.py (the planner's choices), and
gateway-auth-core mirrors the same split.

Because every credential check now runs before any authorization check, three
answers change. An expired key whose team member is over budget gets the 401
expired-key error instead of the budget error, and no team member budget alert
fires. A blocked or unknown virtual key on a master-key-only route gets the
blocked or invalid-key 401 instead of the master-key-only 401. A virtual key on
a master-key-only route against a proxy without a database gets the 400 no-db
error instead of the master-key-only 401.
This commit is contained in:
Yujong Lee 2026-11-10 22:38:27 +09:00
parent 601211a0d5
commit 9a79b2b2a9
30 changed files with 949 additions and 871 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm-rust/crates/gateway-auth-core/AGENTS.md
View file

@ -1,9 +1,9 @@
# gateway-auth-core
The Rust side of `litellm/proxy/auth/core`. It mirrors that package's facets, hosts, inputs, outcomes and answers type for type, so a facet can move from `PythonCore` into Rust without changing what the proxy hands in or gets back. Read `litellm/proxy/auth/core/AGENTS.md` first, its rules apply here unchanged
The Rust side of `litellm/proxy/auth/core`. It mirrors that package's facets, hosts, inputs, identities, outcomes, validators and answers type for type, so a facet can move from `PythonCore` into Rust without changing what the proxy hands in or gets back. Read `litellm/proxy/auth/core/AGENTS.md` first, its rules apply here unchanged
The crate is pure. It depends on no IO, runtime, HTTP, database or Python crate, and a facet reaches outside itself only through its host trait. Each host trait has its own `Error` type, so a pyo3-backed host and a native host can fail in their own terms while the facet code stays the same
Each Python protocol is a trait here. A core hands out its facets through accessors whose types are associated types on `AuthCore`, the same composition as the Python accessors. Host methods that do IO are `-> impl Future + Send` so facets run under Axum and `tokio::spawn` without `#[async_trait]` boxing. Values a host returns as untyped JSON in Python, such as introspection fields, are `serde_json::Value` here
Each Python protocol is a trait here. A core hands out its facets through accessors whose types are associated types on `AuthCore`, the same composition as the Python accessors. Host methods that do IO are `-> impl Future + Send` so facets run under Axum and `tokio::spawn` without `#[async_trait]` boxing. Values a host returns as untyped JSON in Python, such as introspection fields, are `serde_json::Value` here. A Python union of an identity and a denial is a `Result` here, as in `KeyAuthentication`
`NativeCore` implements `AuthCore` with every facet `Unported`, which panics if called, the same as `RustCore` in Python. The crate exists so the types stay in step with Python. A facet leaves `Unported` only together with the Python change that routes the proxy to it, and both sides then share one test suite of cases

15
litellm-rust/crates/gateway-auth-core/src/answers.rs
View file

@ -81,18 +81,3 @@ pub struct Introspection {
pub role: Value,
pub team_id: Value,
}
#[derive(Clone, Debug, PartialEq)]
pub struct TeamFacts {
pub blocked: bool,
pub models: Vec<String>,
pub spend: f64,
pub max_budget: Option<f64>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct OrganizationFacts {
pub models: Vec<String>,
pub spend: f64,
pub max_budget: Option<f64>,
}

23
litellm-rust/crates/gateway-auth-core/src/facets.rs
View file

@ -5,15 +5,17 @@ use crate::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
},
identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication,
KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial,
OAuth2ProxyIdentity, VirtualKey,
},
inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings,
OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
},
outcome::{
Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan,
ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity,
Outcome, Plan,
},
outcome::{ModelBudgetResult, Outcome},
validators::{CredentialPlan, KeyPlan, Plan},
};
pub trait KeyCore: Sync {
@ -23,14 +25,13 @@ pub trait KeyCore: Sync {
facts: &RequestFacts,
settings: &KeySettings,
now: SystemTime,
) -> impl Future<Output = Result<Outcome, H::Error>> + Send;
fn check<H: KeyHost>(
) -> impl Future<Output = Result<KeyAuthentication, H::Error>> + Send;
fn verify<H: KeyHost>(
&self,
host: &H,
facts: &RequestFacts,
settings: &KeySettings,
now: SystemTime,
) -> impl Future<Output = Result<Outcome, H::Error>> + Send;
) -> impl Future<Output = Result<Result<VirtualKey, KeyDenial>, H::Error>> + Send;
}
pub trait ModelBudgetCore: Sync {
@ -71,7 +72,7 @@ pub trait JwtCore: Sync {
facts: &RequestFacts,
settings: &JwtSettings,
now: SystemTime,
) -> impl Future<Output = Result<Result<Identity, Denial>, H::Error>> + Send;
) -> impl Future<Output = Result<Result<Identity, AuthDenial>, H::Error>> + Send;
}
pub trait OAuth2Core: Sync {
@ -100,7 +101,7 @@ pub trait AuthorizeCore: Sync {
fn authorize<H: AuthorizeHost>(
&self,
host: &H,
identity: &Identity,
identity: &KeyIdentity,
facts: &RequestFacts,
settings: &AuthorizeSettings,
) -> impl Future<Output = Result<Outcome, H::Error>> + Send;

52
litellm-rust/crates/gateway-auth-core/src/host.rs
View file

@ -2,12 +2,12 @@ use std::{error::Error, future::Future};
use crate::{
answers::{
FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OrganizationFacts, OwnerFacts,
PresentedKey, SigningKey, TeamFacts, TeamMember, TeamMemberBudget,
TeamMemberModelBudgetScope,
FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OwnerFacts, PresentedKey,
SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
},
identity::{CustomAuthResult, JwtIdentity},
inputs::RequestFacts,
outcome::{CustomAuthResult, Identity, JwtIdentity, ModelBudgetExceeded},
outcome::ModelBudgetExceeded,
};
pub trait ModelBudgetHost: Sync {
@ -60,9 +60,6 @@ pub trait TeamMemberModelBudgetHost: Sync {
pub trait KeyHost: Sync {
type Error: Error + Send + Sync + 'static;
type ModelBudgets: ModelBudgetHost;
fn model_budgets(&self) -> Self::ModelBudgets;
fn resolve_end_user(
&self,
@ -82,11 +79,19 @@ pub trait KeyHost: Sync {
loaded_from_db: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn inspect_key(&self) -> impl Future<Output = Result<KeyFacts, Self::Error>> + Send;
fn enforce_model_access(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn load_key_owner(
&self,
user_id: &str,
) -> impl Future<Output = Result<Option<OwnerFacts>, Self::Error>> + Send;
}
pub trait AuthorizeHost: Sync {
type Error: Error + Send + Sync + 'static;
type ModelBudgets: ModelBudgetHost;
fn model_budgets(&self) -> Self::ModelBudgets;
fn enforce_model_access(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn check_budget_exemption(&self) -> impl Future<Output = Result<bool, Self::Error>> + Send;
fn read_team_member_budget(
&self,
@ -112,14 +117,6 @@ pub trait KeyHost: Sync {
) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SpendCounter {
Key { hashed_token: String },
Team { team_id: String },
Organization { org_id: String },
TeamMember(TeamMember),
}
pub trait JwtHost: Sync {
type Error: Error + Send + Sync + 'static;
@ -127,7 +124,7 @@ pub trait JwtHost: Sync {
&self,
issuer: &str,
) -> impl Future<Output = Result<Vec<SigningKey>, Self::Error>> + Send;
fn mapped_key(
fn mapped_key_hash(
&self,
identity: &JwtIdentity,
) -> impl Future<Output = Result<Option<String>, Self::Error>> + Send;
@ -155,24 +152,3 @@ pub trait CustomAuthHost: Sync {
facts: &RequestFacts,
) -> impl Future<Output = Result<CustomAuthResult, Self::Error>> + Send;
}
pub trait AuthorizeHost: Sync {
type Error: Error + Send + Sync + 'static;
fn team(
&self,
team_id: &str,
) -> impl Future<Output = Result<Option<TeamFacts>, Self::Error>> + Send;
fn organization(
&self,
org_id: &str,
) -> impl Future<Output = Result<Option<OrganizationFacts>, Self::Error>> + Send;
fn allowed_models(
&self,
identity: &Identity,
) -> impl Future<Output = Result<Vec<String>, Self::Error>> + Send;
fn spend(
&self,
counter: &SpendCounter,
) -> impl Future<Output = Result<f64, Self::Error>> + Send;
}

126
litellm-rust/crates/gateway-auth-core/src/identity.rs Normal file
View file

@ -0,0 +1,126 @@
use std::time::SystemTime;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::answers::KeyFacts;
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum AnonymousReason {
PublicRoute,
PassThroughAuthOff,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Anonymous {
pub reason: AnonymousReason,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct JwtIdentity {
pub user_id: Option<String>,
pub team_id: Option<String>,
pub org_id: Option<String>,
pub end_user_id: Option<String>,
pub role: Option<String>,
pub scopes: Vec<String>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct VirtualKey {
pub key: KeyFacts,
}
#[derive(Clone, Debug, PartialEq)]
pub enum KeyIdentity {
NoMasterKey { token: Option<SecretValue> },
UserAuthRoute,
MasterKey,
CachedAdminKey,
VirtualKey(VirtualKey),
}
#[derive(Clone, Debug, PartialEq)]
pub enum KeyDenial {
MissingKey,
MalformedKey,
NoDatabase,
NotAVirtualKey,
KeyBlocked,
OwnerDeactivated {
user_id: String,
},
ExpiredKey {
expires: SystemTime,
now: SystemTime,
key_label: SecretValue,
},
KeyWithoutToken,
}
pub type KeyAuthentication = Result<KeyIdentity, KeyDenial>;
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct OAuth2Identity {
pub token: String,
pub user_id: Value,
pub role: Value,
pub team_id: Value,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct OAuth2ProxyIdentity {
pub user_id: Option<String>,
pub user_email: Option<String>,
pub team_id: Option<String>,
pub team_alias: Option<String>,
pub org_id: Option<String>,
pub models: Option<Vec<String>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CustomIdentity {
pub user_id: Option<String>,
pub team_id: Option<String>,
pub hashed_token: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum CustomAuthResult {
Identity(CustomIdentity),
ReplacedKey { token: String },
NoIdentity,
}
#[derive(Clone, Debug, PartialEq)]
pub enum Identity {
Key(KeyIdentity),
Jwt(JwtIdentity),
OAuth2(OAuth2Identity),
OAuth2Proxy(OAuth2ProxyIdentity),
Custom(CustomIdentity),
Anonymous(Anonymous),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum OAuth2Denial {
InactiveToken,
IntrospectionNotConfigured,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OAuth2ProxyDenial {
TrustedProxyRangesMissing,
UntrustedProxy { client_ip: Option<String> },
ProxyMappingsMissing,
DisallowedProxyMappings { fields: Vec<String> },
}
#[derive(Clone, Debug, PartialEq)]
pub enum AuthDenial {
Key(KeyDenial),
OAuth2(OAuth2Denial),
OAuth2Proxy(OAuth2ProxyDenial),
}

7
litellm-rust/crates/gateway-auth-core/src/inputs.rs
View file

@ -45,10 +45,7 @@ pub struct PlanSettings {
#[derive(Clone, Debug, PartialEq)]
pub struct KeySettings {
pub master_key: Option<SecretValue>,
pub allow_user_auth: bool,
pub has_db: bool,
pub master_key_only_routes: BTreeSet<String>,
pub proxy_max_budget: f64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
@ -71,6 +68,8 @@ pub struct OAuth2ProxySettings {
#[derive(Clone, Debug, PartialEq)]
pub struct AuthorizeSettings {
pub admin_only_routes: BTreeSet<String>,
pub allow_user_auth: bool,
pub has_db: bool,
pub master_key_only_routes: BTreeSet<String>,
pub proxy_max_budget: f64,
}

22
litellm-rust/crates/gateway-auth-core/src/lib.rs
View file

@ -1,14 +1,15 @@
mod answers;
mod facets;
mod host;
mod identity;
mod inputs;
mod native;
mod outcome;
mod validators;
pub use answers::{
FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope, OrganizationFacts,
OwnerFacts, PresentedKey, SigningKey, TeamFacts, TeamMember, TeamMemberBudget,
TeamMemberModelBudgetScope,
FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope, OwnerFacts,
PresentedKey, SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
};
pub use facets::{
AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore,
@ -16,16 +17,17 @@ pub use facets::{
};
pub use host::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, SpendCounter, TeamMemberModelBudgetHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
};
pub use identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, CustomIdentity, Identity,
JwtIdentity, KeyAuthentication, KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity,
OAuth2ProxyDenial, OAuth2ProxyIdentity, VirtualKey,
};
pub use inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings,
PlanFacts, PlanSettings, RequestFacts, RoutingOverride,
};
pub use native::{NativeCore, Unported};
pub use outcome::{
Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, CustomIdentity, Denial, Grant,
Identity, JwtIdentity, KeyPlan, ModelBudgetExceeded, ModelBudgetResult, OAuth2Denial,
OAuth2Identity, OAuth2Mode, OAuth2ProxyDenial, OAuth2ProxyIdentity, Outcome, Plan,
PremiumFeature, WithinModelBudgets,
};
pub use outcome::{Denial, ModelBudgetExceeded, ModelBudgetResult, Outcome, WithinModelBudgets};
pub use validators::{CredentialPlan, KeyPlan, OAuth2Mode, Plan, PremiumFeature};

23
litellm-rust/crates/gateway-auth-core/src/native.rs
View file

@ -9,15 +9,17 @@ use crate::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
},
identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication,
KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial,
OAuth2ProxyIdentity, VirtualKey,
},
inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings,
OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
},
outcome::{
Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan,
ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity,
Outcome, Plan,
},
outcome::{ModelBudgetResult, Outcome},
validators::{CredentialPlan, KeyPlan, Plan},
};
#[derive(Clone, Copy, Debug, Default)]
@ -78,17 +80,16 @@ impl KeyCore for Unported {
_facts: &RequestFacts,
_settings: &KeySettings,
_now: SystemTime,
) -> Result<Outcome, H::Error> {
) -> Result<KeyAuthentication, H::Error> {
unimplemented!("{RUNS_IN_PYTHON}")
}
async fn check<H: KeyHost>(
async fn verify<H: KeyHost>(
&self,
_host: &H,
_facts: &RequestFacts,
_settings: &KeySettings,
_now: SystemTime,
) -> Result<Outcome, H::Error> {
) -> Result<Result<VirtualKey, KeyDenial>, H::Error> {
unimplemented!("{RUNS_IN_PYTHON}")
}
}
@ -157,7 +158,7 @@ impl JwtCore for Unported {
_facts: &RequestFacts,
_settings: &JwtSettings,
_now: SystemTime,
) -> Result<Result<Identity, Denial>, H::Error> {
) -> Result<Result<Identity, AuthDenial>, H::Error> {
unimplemented!("{RUNS_IN_PYTHON}")
}
}
@ -195,7 +196,7 @@ impl AuthorizeCore for Unported {
async fn authorize<H: AuthorizeHost>(
&self,
_host: &H,
_identity: &Identity,
_identity: &KeyIdentity,
_facts: &RequestFacts,
_settings: &AuthorizeSettings,
) -> Result<Outcome, H::Error> {

152
litellm-rust/crates/gateway-auth-core/src/outcome.rs
View file

@ -1,18 +1,3 @@
use std::time::SystemTime;
use litellm_secret_value::SecretValue;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Debug, PartialEq)]
pub enum Grant {
NoMasterKey { token: Option<SecretValue> },
UserAuthRoute,
MasterKey,
CachedAdminKey,
VirtualKey { budget_throttle_pct: Option<f64> },
}
#[derive(Clone, Debug, PartialEq)]
pub struct ModelBudgetExceeded {
pub message: String,
@ -22,23 +7,15 @@ pub struct ModelBudgetExceeded {
pub entity_id: Option<String>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WithinModelBudgets;
pub type ModelBudgetResult = Result<WithinModelBudgets, ModelBudgetExceeded>;
#[derive(Clone, Debug, PartialEq)]
pub enum Denial {
MissingKey,
MalformedKey,
UserAuthDisabled,
ExpiredKey {
expires: SystemTime,
now: SystemTime,
key_label: SecretValue,
},
MasterKeyOnlyRoute,
NoDatabase,
NotAVirtualKey,
KeyBlocked,
OwnerDeactivated {
user_id: String,
},
TeamMemberBudgetExceeded {
user_id: String,
team_id: String,
@ -49,128 +26,11 @@ pub enum Denial {
spend: f64,
max_budget: f64,
},
KeyWithoutToken,
ModelBudgetExceeded(ModelBudgetExceeded),
}
#[derive(Clone, Debug, PartialEq)]
pub enum Outcome {
Allow(Grant),
Allow { budget_throttle_pct: Option<f64> },
Deny(Denial),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WithinModelBudgets;
pub type ModelBudgetResult = Result<WithinModelBudgets, ModelBudgetExceeded>;
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum AnonymousReason {
PublicRoute,
PassThroughAuthOff,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Anonymous {
pub reason: AnonymousReason,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct JwtIdentity {
pub user_id: Option<String>,
pub team_id: Option<String>,
pub org_id: Option<String>,
pub end_user_id: Option<String>,
pub role: Option<String>,
pub scopes: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct OAuth2Identity {
pub token: String,
pub user_id: Value,
pub role: Value,
pub team_id: Value,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
pub struct OAuth2ProxyIdentity {
pub user_id: Option<String>,
pub user_email: Option<String>,
pub team_id: Option<String>,
pub team_alias: Option<String>,
pub org_id: Option<String>,
pub models: Option<Vec<String>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CustomIdentity {
pub user_id: Option<String>,
pub team_id: Option<String>,
pub hashed_token: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum CustomAuthResult {
Identity(CustomIdentity),
ReplacedKey { token: String },
NoIdentity,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Identity {
MasterKey,
Key { hashed_token: String },
Jwt(JwtIdentity),
OAuth2(OAuth2Identity),
OAuth2Proxy(OAuth2ProxyIdentity),
Custom(CustomIdentity),
Anonymous(Anonymous),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum OAuth2Denial {
InactiveToken,
IntrospectionNotConfigured,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OAuth2ProxyDenial {
TrustedProxyRangesMissing,
UntrustedProxy { client_ip: Option<String> },
ProxyMappingsMissing,
DisallowedProxyMappings { fields: Vec<String> },
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum OAuth2Mode {
Introspection,
Proxy,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum PremiumFeature {
OAuth2,
Jwt,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KeyPlan {
Key { pass_through_key: bool },
Anonymous(AnonymousReason),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlan {
Key(KeyPlan),
Jwt,
OAuth2(OAuth2Mode),
PremiumRequired(PremiumFeature),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Plan {
Credential(CredentialPlan),
Custom,
}

33
litellm-rust/crates/gateway-auth-core/src/validators.rs Normal file
View file

@ -0,0 +1,33 @@
use crate::identity::AnonymousReason;
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum OAuth2Mode {
Introspection,
Proxy,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum PremiumFeature {
OAuth2,
Jwt,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KeyPlan {
Key { pass_through_key: bool },
Anonymous(AnonymousReason),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlan {
Key(KeyPlan),
Jwt,
OAuth2(OAuth2Mode),
PremiumRequired(PremiumFeature),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Plan {
Credential(CredentialPlan),
Custom,
}

15
litellm/proxy/auth/auth_denial.py
View file

@ -7,19 +7,22 @@ from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import key_budget_exceeded_error
from litellm.proxy.auth.auth_utils import abbreviate_api_key
from litellm.proxy.auth.core.outcome import (
Denial,
from litellm.proxy.auth.core.identity import (
ExpiredKey,
KeyBlocked,
KeyBudgetExceeded,
KeyDenial,
KeyWithoutToken,
MalformedKey,
MasterKeyOnlyRoute,
MissingKey,
ModelBudgetExceeded,
NoDatabase,
NotAVirtualKey,
OwnerDeactivated,
)
from litellm.proxy.auth.core.outcome import (
Denial,
KeyBudgetExceeded,
MasterKeyOnlyRoute,
ModelBudgetExceeded,
TeamMemberBudgetExceeded,
UserAuthDisabled,
)
@ -45,7 +48,7 @@ def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException:
def denial_exception(
denial: Denial,
denial: KeyDenial | Denial,
token: str | None,
route: str,
jwt_auth_enabled: bool,

12
litellm/proxy/auth/core/AGENTS.md
View file

@ -10,19 +10,19 @@ Modules here import only the standard library and each other. A core does no IO
`core.py` splits the decision into facets, each a protocol with its own host. A core hands each facet out through an accessor, the way `RouterCore` hands out `registry()`. Nothing inherits from another facet.
`AuthCore` has one accessor per facet. `plan()` picks the credential path a request takes, in the proxy's order: custom auth, public routes, OAuth2 introspection, OAuth2 proxy headers, JWT, unauthenticated pass-through routes, then keys. A path that can hand the request on (custom auth issuing a key, a JWT mapped to a virtual key) resumes planning through `after_custom_auth()` or `after_jwt()`, whose narrower return types exclude the paths already taken. `anonymous()`, `oauth2()`, `custom()` and `jwt()` turn a credential into an identity. `keys()` authenticates a key credential (master key, virtual key, UI session, cached admin key) and re-checks a key the proxy resolved itself. `model_budgets()` checks per-model budgets for the key and user, the user alone, or a team member. `authorize()` decides whether an identity may make the request.
`AuthCore` has one accessor per facet. `plan()` picks the credential path a request takes, in the proxy's order: custom auth, public routes, OAuth2 introspection, OAuth2 proxy headers, JWT, unauthenticated pass-through routes, then keys. A path that can hand the request on (custom auth issuing a key, a JWT mapped to a virtual key) resumes planning through `after_custom_auth()` or `after_jwt()`, whose narrower return types exclude the paths already taken. `anonymous()`, `oauth2()`, `custom()` and `jwt()` turn a credential into an identity. `keys()` authenticates a key credential (master key, virtual key, UI session, cached admin key) and verifies a key the proxy resolved itself. `authorize()` decides whether a key identity may make the request: master-key-only routes, model access and the key, team member and per-model budgets. `model_budgets()` checks per-model budgets for the key and user, the user alone, or a team member.
`litellm/proxy/auth/selected_core.py` picks the core the proxy runs, `PythonCore`. `RustCore` declares every accessor and implements none. `PythonCore` does not implement `jwt()` or `authorize()` yet, so their signatures are provisional and the proxy still runs that logic itself.
`litellm/proxy/auth/selected_core.py` picks the core the proxy runs, `PythonCore`. `RustCore` declares every accessor and implements none. `PythonCore` does not implement `jwt()` yet, so its signature is provisional and the proxy still validates JWTs itself. `authorize()` covers key identities only, and every identity still passes through the proxy's common checks afterwards.
Each facet returns only its own results, and the proxy converts each one separately. `KeyCore` returns an `Outcome` (`Allow` with a key `Grant`, or `Deny` with a `Denial`), `ModelBudgetCore` returns `WithinModelBudgets` or `ModelBudgetExceeded`, and the planner returns a validator or `PremiumRequired`.
Authentication and authorization are separate steps with separate results. An authenticating facet returns an identity or the reason the credential is unusable (`keys()` returns a `KeyAuthentication`). `authorize()` takes that identity and returns an `Outcome`, `Allow` or `Deny` with a `Denial`. A check about the credential itself (missing, malformed, blocked, expired, owner deactivated) belongs in authentication, and a check about what the identity may do (routes, models, budgets) belongs in `authorize()`, so a credential that is unusable is always refused as such before any budget is read or alert sent. `ModelBudgetCore` returns `WithinModelBudgets` or `ModelBudgetExceeded`, and the planner returns a validator or `PremiumRequired`.
## Hosts carry the IO
`host.py` lists every read, write and alert a facet may ask for, as typed methods, one host per kind of IO a facet needs (OAuth2 introspection calls out over HTTP while proxy-header mode only reads the request, so each has its own host). A host performs IO and returns data. Anything that compares, chooses or denies belongs in the core. When a decision needs data the core doesn't have, add a host method that returns that data instead of letting the host decide. When one check runs another, its host hands out the other host through an accessor, as `KeyHost.model_budgets()` does.
`host.py` lists every read, write and alert a facet may ask for, as typed methods, one host per kind of IO a facet needs (OAuth2 introspection calls out over HTTP while proxy-header mode only reads the request, so each has its own host). A host performs IO and returns data. Anything that compares, chooses or denies belongs in the core. When a decision needs data the core doesn't have, add a host method that returns that data instead of letting the host decide. When one check runs another, its host hands out the other host through an accessor, as `AuthorizeHost.model_budgets()` does.
## Model the data, not flags
The values that cross the boundary live in one module per direction. `inputs.py` holds what a caller passes in (`RequestFacts`, `PlanFacts` and the settings each facet reads), `outcome.py` holds what a facet returns (grants, denials, identities and validators), and `answers.py` holds what host methods return. Each is a frozen dataclass, and each choice is a union of them. Express "absent" as `None` on a typed value (a key without a team membership has `team_member=None`) instead of a boolean next to optional fields.
The values that cross the boundary live in one module per direction. `inputs.py` holds what a caller passes in (`RequestFacts`, `PlanFacts` and the settings each facet reads), `identity.py` holds what authentication returns (identities and why a credential was refused), `outcome.py` holds the authorization verdict, `validators.py` holds the planner's choices, and `answers.py` holds what host methods return. Each is a frozen dataclass, and each choice is a union of them. Express "absent" as `None` on a typed value (a key without a team membership has `team_member=None`) instead of a boolean next to optional fields.
## Order is behavior
@ -30,7 +30,7 @@ The order in which a core calls host methods is part of its contract. It decides
## Projection stays outside
Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exception types. The proxy turns each grant into its token object and maps every `Denial` to the public error contract in one place, `litellm/proxy/auth/auth_denial.py`. A new denial needs a case there.
Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exception types. The proxy turns each identity into its token object and maps every `KeyDenial` and `Denial` to the public error contract in one place, `litellm/proxy/auth/auth_denial.py`. A new denial needs a case there.
## Credentials

15
litellm/proxy/auth/core/answers.py
View file

@ -88,18 +88,3 @@ class Introspection:
user_id: object
role: object
team_id: object
@dataclass(frozen=True, slots=True)
class TeamFacts:
blocked: bool
models: tuple[str, ...]
spend: float
max_budget: float | None
@dataclass(frozen=True, slots=True)
class OrganizationFacts:
models: tuple[str, ...]
spend: float
max_budget: float | None

41
litellm/proxy/auth/core/core.py
View file

@ -11,6 +11,21 @@ from litellm.proxy.auth.core.host import (
OAuth2ProxyHost,
TeamMemberModelBudgetHost,
)
from litellm.proxy.auth.core.identity import (
Anonymous,
AnonymousReason,
AuthDenial,
CustomAuthResult,
Identity,
KeyAuthentication,
KeyDenial,
KeyIdentity,
OAuth2Denial,
OAuth2Identity,
OAuth2ProxyDenial,
OAuth2ProxyIdentity,
VirtualKey,
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
JwtSettings,
@ -21,30 +36,16 @@ from litellm.proxy.auth.core.inputs import (
PlanSettings,
RequestFacts,
)
from litellm.proxy.auth.core.outcome import (
Anonymous,
AnonymousReason,
CredentialPlan,
CustomAuthResult,
Denial,
Identity,
KeyPlan,
ModelBudgetResult,
OAuth2Denial,
OAuth2Identity,
OAuth2ProxyDenial,
OAuth2ProxyIdentity,
Outcome,
Plan,
)
from litellm.proxy.auth.core.outcome import ModelBudgetResult, Outcome
from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan
class KeyCore(Protocol):
async def authenticate(
self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime
) -> Outcome: ...
) -> KeyAuthentication: ...
async def check(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome: ...
async def verify(self, host: KeyHost, facts: RequestFacts, now: datetime) -> VirtualKey | KeyDenial: ...
class ModelBudgetCore(Protocol):
@ -72,7 +73,7 @@ class AnonymousCore(Protocol):
class JwtCore(Protocol):
async def authenticate(
self, host: JwtHost, facts: RequestFacts, settings: JwtSettings, now: datetime
) -> Identity | Denial: ...
) -> Identity | AuthDenial: ...
class OAuth2Core(Protocol):
@ -91,7 +92,7 @@ class CustomAuthCore(Protocol):
class AuthorizeCore(Protocol):
async def authorize(
self, host: AuthorizeHost, identity: Identity, facts: RequestFacts, settings: AuthorizeSettings
self, host: AuthorizeHost, identity: KeyIdentity, facts: RequestFacts, settings: AuthorizeSettings
) -> Outcome: ...

53
litellm/proxy/auth/core/host.py
View file

@ -1,22 +1,20 @@
from dataclasses import dataclass
from typing import Protocol, TypeAlias
from typing import Protocol
from litellm.proxy.auth.core.answers import (
FallbackTarget,
Introspection,
KeyFacts,
ModelBudgetScope,
OrganizationFacts,
OwnerFacts,
PresentedKey,
SigningKey,
TeamFacts,
TeamMember,
TeamMemberBudget,
TeamMemberModelBudgetScope,
)
from litellm.proxy.auth.core.identity import CustomAuthResult, JwtIdentity
from litellm.proxy.auth.core.inputs import RequestFacts
from litellm.proxy.auth.core.outcome import CustomAuthResult, Identity, JwtIdentity, KeyIdentity, ModelBudgetExceeded
from litellm.proxy.auth.core.outcome import ModelBudgetExceeded
class ModelBudgetHost(Protocol):
@ -44,8 +42,6 @@ class TeamMemberModelBudgetHost(Protocol):
class KeyHost(Protocol):
def model_budgets(self) -> ModelBudgetHost: ...
async def resolve_end_user(self, prefetch_key: bool) -> None: ...
async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ...
@ -58,10 +54,14 @@ class KeyHost(Protocol):
async def inspect_key(self) -> KeyFacts: ...
async def enforce_model_access(self) -> None: ...
async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ...
class AuthorizeHost(Protocol):
def model_budgets(self) -> ModelBudgetHost: ...
async def enforce_model_access(self) -> None: ...
async def check_budget_exemption(self) -> bool: ...
async def read_team_member_budget(self, member: TeamMember) -> TeamMemberBudget | None: ...
@ -83,33 +83,10 @@ class KeyHost(Protocol):
async def alert_proxy_budget(self, spend: float) -> None: ...
@dataclass(frozen=True, slots=True)
class KeySpend:
hashed_token: str
@dataclass(frozen=True, slots=True)
class TeamSpend:
team_id: str
@dataclass(frozen=True, slots=True)
class OrganizationSpend:
org_id: str
@dataclass(frozen=True, slots=True)
class TeamMemberSpend:
member: TeamMember
SpendCounter: TypeAlias = KeySpend | TeamSpend | OrganizationSpend | TeamMemberSpend
class JwtHost(Protocol):
async def signing_keys(self, issuer: str) -> tuple[SigningKey, ...]: ...
async def mapped_key(self, identity: JwtIdentity) -> KeyIdentity | None: ...
async def mapped_key_hash(self, identity: JwtIdentity) -> str | None: ...
class OAuth2IntrospectionHost(Protocol):
@ -124,13 +101,3 @@ class OAuth2ProxyHost(Protocol):
class CustomAuthHost(Protocol):
async def custom_auth(self, facts: RequestFacts) -> CustomAuthResult: ...
class AuthorizeHost(Protocol):
async def team(self, team_id: str) -> TeamFacts | None: ...
async def organization(self, org_id: str) -> OrganizationFacts | None: ...
async def allowed_models(self, identity: Identity) -> frozenset[str]: ...
async def spend(self, counter: SpendCounter) -> float: ...

191
litellm/proxy/auth/core/identity.py Normal file
View file

@ -0,0 +1,191 @@
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import TypeAlias
from litellm.proxy.auth.core.answers import KeyFacts
class AnonymousReason(Enum):
PUBLIC_ROUTE = "public_route"
PASS_THROUGH_AUTH_OFF = "pass_through_auth_off"
@dataclass(frozen=True, slots=True)
class Anonymous:
reason: AnonymousReason
@dataclass(frozen=True, slots=True)
class NoMasterKey:
token: str | None
@dataclass(frozen=True, slots=True)
class UserAuthRoute:
pass
@dataclass(frozen=True, slots=True)
class MasterKey:
pass
@dataclass(frozen=True, slots=True)
class CachedAdminKey:
pass
@dataclass(frozen=True, slots=True)
class VirtualKey:
key: KeyFacts
KeyIdentity: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey
@dataclass(frozen=True, slots=True)
class JwtIdentity:
user_id: str | None
team_id: str | None
org_id: str | None
end_user_id: str | None
role: str | None
scopes: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class OAuth2Identity:
token: str
user_id: object
role: object
team_id: object
@dataclass(frozen=True, slots=True)
class OAuth2ProxyIdentity:
user_id: str | None
user_email: str | None
team_id: str | None
team_alias: str | None
org_id: str | None
models: tuple[str, ...] | None
@dataclass(frozen=True, slots=True)
class CustomIdentity:
user_id: str | None
team_id: str | None
hashed_token: str | None
Identity: TypeAlias = KeyIdentity | JwtIdentity | OAuth2Identity | OAuth2ProxyIdentity | CustomIdentity | Anonymous
@dataclass(frozen=True, slots=True)
class ReplacedKey:
token: str
@dataclass(frozen=True, slots=True)
class NoCustomIdentity:
pass
CustomAuthResult: TypeAlias = CustomIdentity | ReplacedKey | NoCustomIdentity
@dataclass(frozen=True, slots=True)
class MissingKey:
pass
@dataclass(frozen=True, slots=True)
class MalformedKey:
pass
@dataclass(frozen=True, slots=True)
class NoDatabase:
pass
@dataclass(frozen=True, slots=True)
class NotAVirtualKey:
pass
@dataclass(frozen=True, slots=True)
class KeyBlocked:
pass
@dataclass(frozen=True, slots=True)
class OwnerDeactivated:
user_id: str
@dataclass(frozen=True, slots=True)
class ExpiredKey:
expires: datetime
now: datetime
key_label: str
@dataclass(frozen=True, slots=True)
class KeyWithoutToken:
pass
KeyDenial: TypeAlias = (
MissingKey
| MalformedKey
| NoDatabase
| NotAVirtualKey
| KeyBlocked
| OwnerDeactivated
| ExpiredKey
| KeyWithoutToken
)
KeyAuthentication: TypeAlias = KeyIdentity | KeyDenial
@dataclass(frozen=True, slots=True)
class InactiveToken:
pass
@dataclass(frozen=True, slots=True)
class IntrospectionNotConfigured:
pass
OAuth2Denial: TypeAlias = InactiveToken | IntrospectionNotConfigured
@dataclass(frozen=True, slots=True)
class TrustedProxyRangesMissing:
pass
@dataclass(frozen=True, slots=True)
class UntrustedProxy:
client_ip: str | None
@dataclass(frozen=True, slots=True)
class ProxyMappingsMissing:
pass
@dataclass(frozen=True, slots=True)
class DisallowedProxyMappings:
fields: tuple[str, ...]
OAuth2ProxyDenial: TypeAlias = (
TrustedProxyRangesMissing | UntrustedProxy | ProxyMappingsMissing | DisallowedProxyMappings
)
AuthDenial: TypeAlias = KeyDenial | OAuth2Denial | OAuth2ProxyDenial

7
litellm/proxy/auth/core/inputs.py
View file

@ -43,10 +43,7 @@ class PlanSettings:
@dataclass(frozen=True, slots=True)
class KeySettings:
master_key: str | None
allow_user_auth: bool
has_db: bool
master_key_only_routes: frozenset[str]
proxy_max_budget: float
@dataclass(frozen=True, slots=True)
@ -72,5 +69,7 @@ class OAuth2ProxySettings:
@dataclass(frozen=True, slots=True)
class AuthorizeSettings:
admin_only_routes: frozenset[str]
allow_user_auth: bool
has_db: bool
master_key_only_routes: frozenset[str]
proxy_max_budget: float

4
litellm/proxy/auth/core/model_budgets.py
View file

@ -3,7 +3,9 @@ from typing import Final
from litellm.proxy.auth.core.answers import FallbackTarget, ModelBudgetScope, TeamMemberModelBudgetScope
from litellm.proxy.auth.core.host import ModelBudgetHost, TeamMemberModelBudgetHost
from litellm.proxy.auth.core.outcome import ModelBudgetExceeded
from litellm.proxy.auth.core.outcome import (
ModelBudgetExceeded,
)
async def check_key_and_user_model_budgets(host: ModelBudgetHost) -> ModelBudgetExceeded | None:

6
litellm/proxy/auth/core/oauth2.py
View file

@ -1,9 +1,7 @@
from typing import Final
from litellm.proxy.auth.core.host import OAuth2IntrospectionHost, OAuth2ProxyHost
from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings
from litellm.proxy.auth.core.network import ip_in_networks
from litellm.proxy.auth.core.outcome import (
from litellm.proxy.auth.core.identity import (
DisallowedProxyMappings,
InactiveToken,
IntrospectionNotConfigured,
@ -15,6 +13,8 @@ from litellm.proxy.auth.core.outcome import (
TrustedProxyRangesMissing,
UntrustedProxy,
)
from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings
from litellm.proxy.auth.core.network import ip_in_networks
# Header trust asserts identity only. Mapping any other UserAPIKeyAuth field (role, budgets, limits, permissions)
# would let a caller forge its own enforcement by sending the matching header, and an allowlist stays safe as

260
litellm/proxy/auth/core/outcome.py
View file

@ -1,89 +1,17 @@
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import TypeAlias
@dataclass(frozen=True, slots=True)
class NoMasterKey:
token: str | None
@dataclass(frozen=True, slots=True)
class UserAuthRoute:
pass
@dataclass(frozen=True, slots=True)
class MasterKey:
pass
@dataclass(frozen=True, slots=True)
class CachedAdminKey:
pass
@dataclass(frozen=True, slots=True)
class VirtualKey:
budget_throttle_pct: float | None
@dataclass(frozen=True, slots=True)
class WithinModelBudgets:
pass
Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey
@dataclass(frozen=True, slots=True)
class MissingKey:
pass
@dataclass(frozen=True, slots=True)
class MalformedKey:
pass
@dataclass(frozen=True, slots=True)
class UserAuthDisabled:
pass
@dataclass(frozen=True, slots=True)
class ExpiredKey:
expires: datetime
now: datetime
key_label: str
@dataclass(frozen=True, slots=True)
class MasterKeyOnlyRoute:
pass
@dataclass(frozen=True, slots=True)
class NoDatabase:
pass
@dataclass(frozen=True, slots=True)
class NotAVirtualKey:
pass
@dataclass(frozen=True, slots=True)
class KeyBlocked:
pass
@dataclass(frozen=True, slots=True)
class OwnerDeactivated:
user_id: str
@dataclass(frozen=True, slots=True)
class TeamMemberBudgetExceeded:
user_id: str
@ -98,11 +26,6 @@ class KeyBudgetExceeded:
max_budget: float
@dataclass(frozen=True, slots=True)
class KeyWithoutToken:
pass
@dataclass(frozen=True, slots=True)
class ModelBudgetExceeded:
message: str
@ -112,26 +35,21 @@ class ModelBudgetExceeded:
entity_id: str | None
@dataclass(frozen=True, slots=True)
class WithinModelBudgets:
pass
ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded
Denial: TypeAlias = (
MissingKey
| MalformedKey
| UserAuthDisabled
| ExpiredKey
| MasterKeyOnlyRoute
| NoDatabase
| NotAVirtualKey
| KeyBlocked
| OwnerDeactivated
| TeamMemberBudgetExceeded
| KeyBudgetExceeded
| KeyWithoutToken
| ModelBudgetExceeded
UserAuthDisabled | MasterKeyOnlyRoute | TeamMemberBudgetExceeded | KeyBudgetExceeded | ModelBudgetExceeded
)
@dataclass(frozen=True, slots=True)
class Allow:
grant: Grant
budget_throttle_pct: float | None
@dataclass(frozen=True, slots=True)
@ -140,163 +58,3 @@ class Deny:
Outcome: TypeAlias = Allow | Deny
ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded
class AnonymousReason(Enum):
PUBLIC_ROUTE = "public_route"
PASS_THROUGH_AUTH_OFF = "pass_through_auth_off"
@dataclass(frozen=True, slots=True)
class MasterKeyIdentity:
pass
@dataclass(frozen=True, slots=True)
class KeyIdentity:
hashed_token: str
@dataclass(frozen=True, slots=True)
class JwtIdentity:
user_id: str | None
team_id: str | None
org_id: str | None
end_user_id: str | None
role: str | None
scopes: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class OAuth2Identity:
token: str
user_id: object
role: object
team_id: object
@dataclass(frozen=True, slots=True)
class OAuth2ProxyIdentity:
user_id: str | None
user_email: str | None
team_id: str | None
team_alias: str | None
org_id: str | None
models: tuple[str, ...] | None
@dataclass(frozen=True, slots=True)
class CustomIdentity:
user_id: str | None
team_id: str | None
hashed_token: str | None
@dataclass(frozen=True, slots=True)
class ReplacedKey:
token: str
@dataclass(frozen=True, slots=True)
class NoCustomIdentity:
pass
CustomAuthResult: TypeAlias = CustomIdentity | ReplacedKey | NoCustomIdentity
@dataclass(frozen=True, slots=True)
class Anonymous:
reason: AnonymousReason
@dataclass(frozen=True, slots=True)
class InactiveToken:
pass
@dataclass(frozen=True, slots=True)
class IntrospectionNotConfigured:
pass
OAuth2Denial: TypeAlias = InactiveToken | IntrospectionNotConfigured
@dataclass(frozen=True, slots=True)
class TrustedProxyRangesMissing:
pass
@dataclass(frozen=True, slots=True)
class UntrustedProxy:
client_ip: str | None
@dataclass(frozen=True, slots=True)
class ProxyMappingsMissing:
pass
@dataclass(frozen=True, slots=True)
class DisallowedProxyMappings:
fields: tuple[str, ...]
OAuth2ProxyDenial: TypeAlias = (
TrustedProxyRangesMissing | UntrustedProxy | ProxyMappingsMissing | DisallowedProxyMappings
)
Identity: TypeAlias = (
MasterKeyIdentity | KeyIdentity | JwtIdentity | OAuth2Identity | OAuth2ProxyIdentity | CustomIdentity | Anonymous
)
class OAuth2Mode(Enum):
INTROSPECTION = "introspection"
PROXY = "proxy"
@dataclass(frozen=True, slots=True)
class KeyValidator:
pass_through_key: bool
@dataclass(frozen=True, slots=True)
class JwtValidator:
pass
@dataclass(frozen=True, slots=True)
class OAuth2Validator:
mode: OAuth2Mode
@dataclass(frozen=True, slots=True)
class CustomValidator:
pass
@dataclass(frozen=True, slots=True)
class AnonymousValidator:
reason: AnonymousReason
class PremiumFeature(Enum):
OAUTH2 = "oauth2"
JWT = "jwt"
@dataclass(frozen=True, slots=True)
class PremiumRequired:
feature: PremiumFeature
KeyPlan: TypeAlias = KeyValidator | AnonymousValidator
CredentialPlan: TypeAlias = KeyPlan | JwtValidator | OAuth2Validator | PremiumRequired
Plan: TypeAlias = CredentialPlan | CustomValidator

6
litellm/proxy/auth/core/plan.py
View file

@ -3,9 +3,11 @@ from fnmatch import fnmatchcase
from typing import Final
from litellm.proxy.auth.core.credential import is_jwt
from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RoutingOverride
from litellm.proxy.auth.core.outcome import (
from litellm.proxy.auth.core.identity import (
AnonymousReason,
)
from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RoutingOverride
from litellm.proxy.auth.core.validators import (
AnonymousValidator,
CredentialPlan,
CustomValidator,

184
litellm/proxy/auth/core/python_core.py
View file

@ -18,6 +18,7 @@ from litellm.proxy.auth.core.core import (
OAuth2Core,
)
from litellm.proxy.auth.core.host import (
AuthorizeHost,
CustomAuthHost,
KeyHost,
ModelBudgetHost,
@ -25,7 +26,33 @@ from litellm.proxy.auth.core.host import (
OAuth2ProxyHost,
TeamMemberModelBudgetHost,
)
from litellm.proxy.auth.core.identity import (
Anonymous,
AnonymousReason,
CachedAdminKey,
CustomAuthResult,
ExpiredKey,
KeyAuthentication,
KeyBlocked,
KeyDenial,
KeyIdentity,
KeyWithoutToken,
MalformedKey,
MasterKey,
MissingKey,
NoDatabase,
NoMasterKey,
NotAVirtualKey,
OAuth2Denial,
OAuth2Identity,
OAuth2ProxyDenial,
OAuth2ProxyIdentity,
OwnerDeactivated,
UserAuthRoute,
VirtualKey,
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
KeySettings,
OAuth2IntrospectionSettings,
OAuth2ProxySettings,
@ -41,52 +68,22 @@ from litellm.proxy.auth.core.model_budgets import (
)
from litellm.proxy.auth.core.outcome import (
Allow,
Anonymous,
AnonymousReason,
CachedAdminKey,
CredentialPlan,
CustomAuthResult,
Deny,
ExpiredKey,
KeyBlocked,
KeyBudgetExceeded,
KeyPlan,
KeyWithoutToken,
MalformedKey,
MasterKey,
MasterKeyOnlyRoute,
MissingKey,
ModelBudgetExceeded,
ModelBudgetResult,
NoDatabase,
NoMasterKey,
NotAVirtualKey,
OAuth2Denial,
OAuth2Identity,
OAuth2ProxyDenial,
OAuth2ProxyIdentity,
Outcome,
OwnerDeactivated,
Plan,
TeamMemberBudgetExceeded,
UserAuthDisabled,
UserAuthRoute,
VirtualKey,
WithinModelBudgets,
)
from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan
_USER_AUTH_ROUTE: Final = "/user/auth"
_VIRTUAL_KEY_PREFIX: Final = "sk-"
@dataclass(frozen=True, slots=True)
class _KeyCheck:
settings: KeySettings
now: datetime
is_llm_api_route: bool
key_label: str
@dataclass(frozen=True, slots=True)
class _WithinKeyBudget:
budget_throttle_pct: float | None
@ -115,7 +112,7 @@ class PythonCore(AuthCore):
return PythonCustomAuthCore()
def authorize(self) -> AuthorizeCore:
raise NotImplementedError
return PythonAuthorizeCore()
class PythonAuthPlanner(AuthPlanner):
@ -153,24 +150,39 @@ class PythonCustomAuthCore(CustomAuthCore):
class PythonKeyCore(KeyCore):
async def authenticate(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome:
async def authenticate(
self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime
) -> KeyAuthentication:
token: Final = facts.token
if settings.master_key is None:
return Allow(NoMasterKey(token=token))
return NoMasterKey(token=token)
if token is None:
return Deny(MissingKey())
return MissingKey()
if token == "":
return Deny(MalformedKey())
return MalformedKey()
if facts.route == _USER_AUTH_ROUTE:
return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled())
return UserAuthRoute()
if is_master_key(token, settings.master_key):
await host.resolve_end_user(prefetch_key=False)
return Allow(MasterKey())
return MasterKey()
await host.resolve_end_user(prefetch_key=True)
return await _presented_key(host, token, facts, settings, now)
return await _presented_key(host, token, settings, now)
async def check(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome:
return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or ""))
async def verify(self, host: KeyHost, facts: RequestFacts, now: datetime) -> VirtualKey | KeyDenial:
return await _verify_key(host, now, key_label=facts.token or "")
class PythonAuthorizeCore(AuthorizeCore):
async def authorize(
self, host: AuthorizeHost, identity: KeyIdentity, facts: RequestFacts, settings: AuthorizeSettings
) -> Outcome:
match identity:
case UserAuthRoute():
return Allow(budget_throttle_pct=None) if settings.allow_user_auth else Deny(UserAuthDisabled())
case NoMasterKey() | MasterKey() | CachedAdminKey():
return Allow(budget_throttle_pct=None)
case VirtualKey(key=key):
return await _authorize_virtual_key(host, key, facts, settings)
class PythonModelBudgetCore(ModelBudgetCore):
@ -188,56 +200,37 @@ def _within_model_budgets(exceeded: ModelBudgetExceeded | None) -> ModelBudgetRe
return WithinModelBudgets() if exceeded is None else exceeded
async def _presented_key(
host: KeyHost, token: str, facts: RequestFacts, settings: KeySettings, now: datetime
) -> Outcome:
async def _presented_key(host: KeyHost, token: str, settings: KeySettings, now: datetime) -> KeyAuthentication:
hashed_token: Final = hash_token(token)
presented: Final = await host.resolve_presented_key(hashed_token)
if isinstance(presented, AdminKey):
if presented.expires is None or presented.expires >= now:
return Allow(CachedAdminKey())
return CachedAdminKey()
await host.evict_presented_key(hashed_token)
return Deny(ExpiredKey(expires=presented.expires, now=now, key_label=token))
if facts.route in settings.master_key_only_routes:
return Deny(MasterKeyOnlyRoute())
return ExpiredKey(expires=presented.expires, now=now, key_label=token)
if not settings.has_db:
return Deny(NoDatabase())
return NoDatabase()
if presented is not None:
await host.prepare_key(loaded_from_db=False)
return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=token))
return await _verify_key(host, now, key_label=token)
if not token.startswith(_VIRTUAL_KEY_PREFIX):
return Deny(NotAVirtualKey())
return NotAVirtualKey()
await host.load_key(hashed_token)
await host.prepare_key(loaded_from_db=True)
return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=hashed_token))
return await _verify_key(host, now, key_label=hashed_token)
async def _check_key(host: KeyHost, check: _KeyCheck) -> Outcome:
async def _verify_key(host: KeyHost, now: datetime, key_label: str) -> VirtualKey | KeyDenial:
key: Final = await host.inspect_key()
if key.blocked:
return Deny(KeyBlocked())
await host.enforce_model_access()
return KeyBlocked()
if key.user_id is not None and await _owner_deactivated(host, key.user_id):
return Deny(OwnerDeactivated(user_id=key.user_id))
exempt: Final = await host.check_budget_exemption()
if not exempt and check.settings.has_db and key.team_member is not None:
team_member_exceeded: Final = await _check_team_member_budget(host, key.team_member)
if team_member_exceeded is not None:
return Deny(team_member_exceeded)
if key.expires is not None and key.expires < check.now:
return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label))
if exempt:
return await _grant(host, check, key, budget_throttle_pct=None)
await host.send_key_budget_alerts()
key_budget: Final = await _check_key_budget(host, check, key)
if isinstance(key_budget, KeyBudgetExceeded):
return Deny(key_budget)
if key.soft_budget and key.spend >= key.soft_budget:
await host.alert_key_soft_budget()
model_budget_exceeded: Final = await check_key_and_user_model_budgets(host.model_budgets())
if model_budget_exceeded is not None:
return Deny(model_budget_exceeded)
return await _grant(host, check, key, budget_throttle_pct=key_budget.budget_throttle_pct)
return OwnerDeactivated(user_id=key.user_id)
if key.expires is not None and key.expires < now:
return ExpiredKey(expires=key.expires, now=now, key_label=key_label)
if not key.has_token:
return KeyWithoutToken()
return VirtualKey(key=key)
async def _owner_deactivated(host: KeyHost, user_id: str) -> bool:
@ -245,7 +238,32 @@ async def _owner_deactivated(host: KeyHost, user_id: str) -> bool:
return owner is not None and owner.deactivated
async def _check_team_member_budget(host: KeyHost, member: TeamMember) -> TeamMemberBudgetExceeded | None:
async def _authorize_virtual_key(
host: AuthorizeHost, key: KeyFacts, facts: RequestFacts, settings: AuthorizeSettings
) -> Outcome:
if facts.route in settings.master_key_only_routes:
return Deny(MasterKeyOnlyRoute())
await host.enforce_model_access()
exempt: Final = await host.check_budget_exemption()
if not exempt and settings.has_db and key.team_member is not None:
team_member_exceeded: Final = await _check_team_member_budget(host, key.team_member)
if team_member_exceeded is not None:
return Deny(team_member_exceeded)
if exempt:
return await _grant(host, settings, budget_throttle_pct=None)
await host.send_key_budget_alerts()
key_budget: Final = await _check_key_budget(host, key, is_llm_api_route=facts.is_llm_api_route)
if isinstance(key_budget, KeyBudgetExceeded):
return Deny(key_budget)
if key.soft_budget and key.spend >= key.soft_budget:
await host.alert_key_soft_budget()
model_budget_exceeded: Final = await check_key_and_user_model_budgets(host.model_budgets())
if model_budget_exceeded is not None:
return Deny(model_budget_exceeded)
return await _grant(host, settings, budget_throttle_pct=key_budget.budget_throttle_pct)
async def _check_team_member_budget(host: AuthorizeHost, member: TeamMember) -> TeamMemberBudgetExceeded | None:
budget: Final = await host.read_team_member_budget(member)
if budget is None or budget.spend < budget.max_budget:
return None
@ -255,9 +273,11 @@ async def _check_team_member_budget(host: KeyHost, member: TeamMember) -> TeamMe
)
async def _check_key_budget(host: KeyHost, check: _KeyCheck, key: KeyFacts) -> _WithinKeyBudget | KeyBudgetExceeded:
async def _check_key_budget(
host: AuthorizeHost, key: KeyFacts, *, is_llm_api_route: bool
) -> _WithinKeyBudget | KeyBudgetExceeded:
max_budget: Final = key.max_budget
if not check.is_llm_api_route or max_budget is None:
if not is_llm_api_route or max_budget is None:
return _WithinKeyBudget(budget_throttle_pct=None)
spend: Final = await host.read_key_spend(max_budget)
await host.alert_key_budget(spend)
@ -269,12 +289,10 @@ async def _check_key_budget(host: KeyHost, check: _KeyCheck, key: KeyFacts) -> _
return KeyBudgetExceeded(spend=spend, max_budget=max_budget)
async def _grant(host: KeyHost, check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Outcome:
async def _grant(host: AuthorizeHost, settings: AuthorizeSettings, budget_throttle_pct: float | None) -> Outcome:
await host.load_key_context()
if check.settings.proxy_max_budget > 0 and check.settings.has_db:
if settings.proxy_max_budget > 0 and settings.has_db:
proxy_spend: Final = await host.read_proxy_spend()
if proxy_spend is not None:
await host.alert_proxy_budget(proxy_spend)
if not key.has_token:
return Deny(KeyWithoutToken())
return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct))
return Allow(budget_throttle_pct=budget_throttle_pct)

52
litellm/proxy/auth/core/validators.py Normal file
View file

@ -0,0 +1,52 @@
from dataclasses import dataclass
from enum import Enum
from typing import TypeAlias
from litellm.proxy.auth.core.identity import AnonymousReason
class OAuth2Mode(Enum):
INTROSPECTION = "introspection"
PROXY = "proxy"
@dataclass(frozen=True, slots=True)
class KeyValidator:
pass_through_key: bool
@dataclass(frozen=True, slots=True)
class JwtValidator:
pass
@dataclass(frozen=True, slots=True)
class OAuth2Validator:
mode: OAuth2Mode
@dataclass(frozen=True, slots=True)
class CustomValidator:
pass
@dataclass(frozen=True, slots=True)
class AnonymousValidator:
reason: AnonymousReason
class PremiumFeature(Enum):
OAUTH2 = "oauth2"
JWT = "jwt"
@dataclass(frozen=True, slots=True)
class PremiumRequired:
feature: PremiumFeature
KeyPlan: TypeAlias = KeyValidator | AnonymousValidator
CredentialPlan: TypeAlias = KeyPlan | JwtValidator | OAuth2Validator | PremiumRequired
Plan: TypeAlias = CredentialPlan | CustomValidator

6
litellm/proxy/auth/oauth2_check.py
View file

@ -12,8 +12,12 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.core.answers import Introspection
from litellm.proxy.auth.core.identity import (
InactiveToken,
IntrospectionNotConfigured,
OAuth2Identity,
)
from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings
from litellm.proxy.auth.core.outcome import InactiveToken, IntrospectionNotConfigured, OAuth2Identity
from litellm.proxy.auth.selected_core import AUTH_CORE
_TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))

6
litellm/proxy/auth/oauth2_proxy_hook.py
View file

@ -7,15 +7,15 @@ from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.core.inputs import OAuth2ProxySettings
from litellm.proxy.auth.core.oauth2 import ALLOWED_OAUTH2_PROXY_FIELDS
from litellm.proxy.auth.core.outcome import (
from litellm.proxy.auth.core.identity import (
DisallowedProxyMappings,
OAuth2ProxyIdentity,
ProxyMappingsMissing,
TrustedProxyRangesMissing,
UntrustedProxy,
)
from litellm.proxy.auth.core.inputs import OAuth2ProxySettings
from litellm.proxy.auth.core.oauth2 import ALLOWED_OAUTH2_PROXY_FIELDS
from litellm.proxy.auth.network import parse_trusted_proxy_ranges
from litellm.proxy.auth.selected_core import AUTH_CORE
from litellm.proxy.auth.trusted_proxy_utils import (

132
litellm/proxy/auth/user_api_key_auth.py
View file

@ -146,36 +146,49 @@ from litellm.proxy.auth.core.credential import (
select_websocket_api_key,
)
from litellm.proxy.auth.core.host import ModelBudgetHost
from litellm.proxy.auth.core.inputs import KeySettings, PlanFacts, PlanSettings, RequestFacts, RoutingOverride
from litellm.proxy.auth.core.outcome import (
Allow,
from litellm.proxy.auth.core.identity import (
Anonymous,
AnonymousReason,
AnonymousValidator,
CachedAdminKey,
CredentialPlan,
CustomAuthResult,
CustomIdentity,
CustomValidator,
Deny,
Grant,
JwtValidator,
KeyValidator,
KeyAuthentication,
KeyDenial,
KeyIdentity,
MasterKey,
ModelBudgetExceeded,
ModelBudgetResult,
NoCustomIdentity,
NoMasterKey,
OAuth2Mode,
OAuth2Validator,
Outcome,
PremiumFeature,
PremiumRequired,
ReplacedKey,
UserAuthRoute,
VirtualKey,
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
KeySettings,
PlanFacts,
PlanSettings,
RequestFacts,
RoutingOverride,
)
from litellm.proxy.auth.core.outcome import (
Allow,
Deny,
ModelBudgetExceeded,
ModelBudgetResult,
Outcome,
WithinModelBudgets,
)
from litellm.proxy.auth.core.validators import (
AnonymousValidator,
CredentialPlan,
CustomValidator,
JwtValidator,
KeyValidator,
OAuth2Mode,
OAuth2Validator,
PremiumFeature,
PremiumRequired,
)
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context
from litellm.proxy.auth.oauth2_check import Oauth2Handler
@ -2525,9 +2538,12 @@ def _proxy_model_list() -> list[object] | None:
_MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value)
def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> KeySettings:
return KeySettings(
master_key=master_key,
def _key_settings(master_key: str | None, prisma_client: PrismaClient | None) -> KeySettings:
return KeySettings(master_key=master_key, has_db=prisma_client is not None)
def _authorize_settings(allow_user_auth: bool, prisma_client: PrismaClient | None) -> AuthorizeSettings:
return AuthorizeSettings(
allow_user_auth=allow_user_auth,
has_db=prisma_client is not None,
master_key_only_routes=_MASTER_KEY_ONLY_ROUTES,
@ -3272,21 +3288,24 @@ async def user_api_key_auth_builder(
proxy_logging_obj=proxy_logging_obj,
route=route,
)
key_facts: Final = RequestFacts(
route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)
)
try:
outcome: Final = await _AUTH_CORE.keys().authenticate(
host,
RequestFacts(route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)),
_core_settings(
master_key=master_key,
allow_user_auth=general_settings.get("allow_user_auth", False) is True,
prisma_client=prisma_client,
authorized: Final = await _authorize_key(
await _AUTH_CORE.keys().authenticate(
host, key_facts, _key_settings(master_key, prisma_client), datetime.now(timezone.utc)
),
host,
key_facts,
_authorize_settings(
allow_user_auth=general_settings.get("allow_user_auth", False) is True, prisma_client=prisma_client
),
datetime.now(timezone.utc),
)
finally:
valid_token = host.resolved_key # rebind-ok: the exception handler below reports the resolved identity
return await _authenticated(
outcome,
authorized,
host,
admin_user_id=litellm_proxy_admin_name,
route=route,
@ -3422,27 +3441,54 @@ async def _master_key_auth(host: _KeyHost, admin_user_id: str, route: str, start
return with_end_user
@dataclass(frozen=True, slots=True)
class _AuthorizedKey:
identity: KeyIdentity
outcome: Outcome
async def _authorize_key(
authentication: KeyAuthentication, host: _KeyHost, facts: RequestFacts, settings: AuthorizeSettings
) -> _AuthorizedKey | KeyDenial:
match authentication:
case NoMasterKey() | UserAuthRoute() | MasterKey() | CachedAdminKey() | VirtualKey():
return _AuthorizedKey(
authentication, await _AUTH_CORE.authorize().authorize(host, authentication, facts, settings)
)
case _:
return authentication
async def _authenticated(
outcome: Outcome,
authorized: _AuthorizedKey | KeyDenial,
host: _KeyHost,
admin_user_id: str,
route: str,
start_time: datetime,
jwt_auth_enabled: bool,
) -> UserAPIKeyAuth:
match outcome:
if not isinstance(authorized, _AuthorizedKey):
raise denial_exception(
authorized, token=host.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=host.resolved_key
)
match authorized.outcome:
case Deny(denial=denial):
raise denial_exception(
denial, token=host.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=host.resolved_key
)
case Allow(grant=grant):
return await _granted(grant, host, admin_user_id, route, start_time)
case Allow(budget_throttle_pct=budget_throttle_pct):
return await _granted(authorized.identity, budget_throttle_pct, host, admin_user_id, route, start_time)
async def _granted(
grant: Grant, host: _KeyHost, admin_user_id: str, route: str, start_time: datetime
identity: KeyIdentity,
budget_throttle_pct: float | None,
host: _KeyHost,
admin_user_id: str,
route: str,
start_time: datetime,
) -> UserAPIKeyAuth:
match grant:
match identity:
case NoMasterKey(token=token):
return UserAPIKeyAuth(
api_key=token, user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=host.parent_otel_span
@ -3453,10 +3499,10 @@ async def _granted(
return await _master_key_auth(host, admin_user_id, route, start_time)
case CachedAdminKey():
return _with_end_user(host.key, host.end_user, host.parent_otel_span)
case VirtualKey(budget_throttle_pct=budget_throttle_pct):
case VirtualKey():
return await _virtual_key_auth(host, budget_throttle_pct, route, start_time)
case _:
assert_never(grant)
assert_never(identity)
async def _virtual_key_auth(
@ -5001,7 +5047,6 @@ async def _authorize_resolved_key(
) -> UserAPIKeyAuth:
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
master_key,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
@ -5019,14 +5064,15 @@ async def _authorize_resolved_key(
proxy_logging_obj=proxy_logging_obj,
route=route,
)
outcome: Final = await _AUTH_CORE.keys().check(
facts: Final = RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route))
authorized: Final = await _authorize_key(
await _AUTH_CORE.keys().verify(host, facts, start_time),
host,
RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)),
_core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client),
start_time,
facts,
_authorize_settings(allow_user_auth=False, prisma_client=prisma_client),
)
return await _authenticated(
outcome,
authorized,
host,
admin_user_id=litellm_proxy_admin_name,
route=route,

6
tests/unit/proxy/auth/core/test_model_budgets.py
View file

@ -10,7 +10,11 @@ from litellm.proxy.auth.core.answers import (
ModelBudgetScope,
TeamMemberModelBudgetScope,
)
from litellm.proxy.auth.core.outcome import ModelBudgetExceeded, ModelBudgetResult, WithinModelBudgets
from litellm.proxy.auth.core.outcome import (
ModelBudgetExceeded,
ModelBudgetResult,
WithinModelBudgets,
)
from litellm.proxy.auth.core.python_core import PythonCore
Call = tuple[str, str | None]

4
tests/unit/proxy/auth/core/test_oauth2.py
View file

@ -7,8 +7,7 @@ from typing import Final
import pytest
from litellm.proxy.auth.core.answers import Introspection
from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings
from litellm.proxy.auth.core.outcome import (
from litellm.proxy.auth.core.identity import (
DisallowedProxyMappings,
InactiveToken,
IntrospectionNotConfigured,
@ -18,6 +17,7 @@ from litellm.proxy.auth.core.outcome import (
TrustedProxyRangesMissing,
UntrustedProxy,
)
from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings
from litellm.proxy.auth.core.python_core import PythonCore
CORE: Final = PythonCore()

8
tests/unit/proxy/auth/core/test_plan.py
View file

@ -4,9 +4,12 @@ from typing import Final
import pytest
from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RequestFacts, RoutingOverride
from litellm.proxy.auth.core.outcome import (
from litellm.proxy.auth.core.identity import (
AnonymousReason,
)
from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RequestFacts, RoutingOverride
from litellm.proxy.auth.core.python_core import PythonCore
from litellm.proxy.auth.core.validators import (
AnonymousValidator,
CustomValidator,
JwtValidator,
@ -16,7 +19,6 @@ from litellm.proxy.auth.core.outcome import (
PremiumFeature,
PremiumRequired,
)
from litellm.proxy.auth.core.python_core import PythonCore
PLANNER: Final = PythonCore().plan()
JWT: Final = "header.payload.signature"

355
tests/unit/proxy/auth/core/test_python_core.py
View file

@ -1,5 +1,4 @@
import math
from collections.abc import Awaitable
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta, timezone
from typing import Final
@ -16,43 +15,49 @@ from litellm.proxy.auth.core.answers import (
TeamMember,
TeamMemberBudget,
)
from litellm.proxy.auth.core.inputs import RequestFacts, KeySettings
from litellm.proxy.auth.core.outcome import (
CustomIdentity,
NoCustomIdentity,
ReplacedKey,
Allow,
from litellm.proxy.auth.core.identity import (
CachedAdminKey,
Deny,
CustomIdentity,
ExpiredKey,
KeyAuthentication,
KeyBlocked,
KeyBudgetExceeded,
KeyIdentity,
KeyWithoutToken,
MalformedKey,
MasterKey,
MasterKeyOnlyRoute,
MissingKey,
ModelBudgetExceeded,
NoCustomIdentity,
NoDatabase,
NoMasterKey,
NotAVirtualKey,
Outcome,
OwnerDeactivated,
TeamMemberBudgetExceeded,
UserAuthDisabled,
ReplacedKey,
UserAuthRoute,
VirtualKey,
)
from litellm.proxy.auth.core.inputs import AuthorizeSettings, KeySettings, RequestFacts
from litellm.proxy.auth.core.key import hash_token
from litellm.proxy.auth.core.outcome import (
Allow,
Deny,
KeyBudgetExceeded,
MasterKeyOnlyRoute,
ModelBudgetExceeded,
Outcome,
TeamMemberBudgetExceeded,
UserAuthDisabled,
)
from litellm.proxy.auth.core.python_core import PythonCore
MASTER: Final = "sk-master"
MASTER_ONLY_ROUTE: Final = "/global/spend/reset"
NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc)
SETTINGS: Final = KeySettings(
master_key=MASTER,
KEY_SETTINGS: Final = KeySettings(master_key=MASTER, has_db=True)
NO_DB: Final = KeySettings(master_key=MASTER, has_db=False)
AUTHORIZE_SETTINGS: Final = AuthorizeSettings(
allow_user_auth=False,
has_db=True,
master_key_only_routes=frozenset({"/global/spend/reset"}),
master_key_only_routes=frozenset({MASTER_ONLY_ROUTE}),
proxy_max_budget=0.0,
)
PLAIN_KEY: Final = KeyFacts(
@ -66,15 +71,25 @@ PLAIN_KEY: Final = KeyFacts(
team_member=None,
budget_throttle_pct=None,
)
CHECKED: Final = (
"prepare_key",
"inspect_key",
AUTHORIZED: Final = (
"enforce_model_access",
"check_budget_exemption",
"send_key_budget_alerts",
"read_model_budget_scope",
"load_key_context",
)
AUTHORIZE_CALLS: Final = frozenset(
{
*AUTHORIZED,
"read_team_member_budget",
"alert_team_member_budget",
"read_key_spend",
"alert_key_budget",
"alert_key_soft_budget",
"read_proxy_spend",
"alert_proxy_budget",
}
)
NO_MODEL_BUDGETS: Final = ModelBudgetScope(
models=(),
@ -190,167 +205,241 @@ def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route
return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route)
async def _run(decision: Awaitable[Outcome], handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]:
outcome: Final = await decision
return outcome, tuple(name for name, _ in handler.calls)
def _calls(handler: RecordingHandler) -> tuple[str, ...]:
return tuple(name for name, _ in handler.calls)
async def _decide(
async def _authenticate(
token: str | None,
*,
route: str = "/chat/completions",
settings: KeySettings = SETTINGS,
settings: KeySettings = KEY_SETTINGS,
handler: RecordingHandler | None = None,
) -> tuple[KeyAuthentication, tuple[str, ...]]:
host: Final = handler or RecordingHandler()
authentication: Final = await PythonCore().keys().authenticate(host, _facts(token, route), settings, NOW)
return authentication, _calls(host)
async def _authorize(
identity: KeyIdentity,
*,
route: str = "/chat/completions",
is_llm_api_route: bool = True,
settings: AuthorizeSettings = AUTHORIZE_SETTINGS,
handler: RecordingHandler | None = None,
) -> tuple[Outcome, tuple[str, ...]]:
host: Final = handler or RecordingHandler()
return await _run(PythonCore().keys().authenticate(host, _facts(token, route), settings, NOW), host)
facts: Final = _facts("hash", route, is_llm_api_route)
outcome: Final = await PythonCore().authorize().authorize(host, identity, facts, settings)
return outcome, _calls(host)
async def _check(
async def _authorize_key(
key: KeyFacts,
*,
is_llm_api_route: bool = True,
settings: KeySettings = SETTINGS,
owner: OwnerFacts | None = None,
exempt: bool = False,
team_member_budget: TeamMemberBudget | None = None,
key_spend: float = 0.0,
model_budget_scope: ModelBudgetScope = NO_MODEL_BUDGETS,
user_model_budget_exceeded: ModelBudgetExceeded | None = None,
) -> tuple[Outcome, tuple[str, ...]]:
handler: Final = RecordingHandler(
key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend
key=key,
exempt=exempt,
team_member_budget=team_member_budget,
key_spend=key_spend,
model_budget_scope=model_budget_scope,
user_model_budget_exceeded=user_model_budget_exceeded,
)
facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route)
return await _run(PythonCore().keys().check(handler, facts, settings, NOW), handler)
return await _authorize(VirtualKey(key=key), handler=handler)
def _allowed(budget_throttle_pct: float | None = None) -> Outcome:
return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct))
return Allow(budget_throttle_pct=budget_throttle_pct)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "token", "settings", "expected"),
[
("/chat/completions", None, KeySettings(None, False, True, frozenset(), 0.0), Allow(NoMasterKey(token=None))),
(
"/chat/completions",
"sk-x",
KeySettings(None, False, False, frozenset(), 0.0),
Allow(NoMasterKey(token="sk-x")),
),
("/chat/completions", None, SETTINGS, Deny(MissingKey())),
("/chat/completions", "", SETTINGS, Deny(MalformedKey())),
("/user/auth", "", KeySettings(MASTER, True, True, frozenset(), 0.0), Deny(MalformedKey())),
("/user/auth", "sk-x", KeySettings(MASTER, True, True, frozenset(), 0.0), Allow(UserAuthRoute())),
("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())),
("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())),
("/chat/completions", None, KeySettings(None, True), NoMasterKey(token=None)),
("/chat/completions", "sk-x", KeySettings(None, False), NoMasterKey(token="sk-x")),
("/chat/completions", None, KEY_SETTINGS, MissingKey()),
("/chat/completions", "", KEY_SETTINGS, MalformedKey()),
("/user/auth", "", KEY_SETTINGS, MalformedKey()),
("/user/auth", "sk-x", KEY_SETTINGS, UserAuthRoute()),
("/user/auth", MASTER, KEY_SETTINGS, UserAuthRoute()),
],
)
async def test_credential_decisions_do_no_io(route: str, token: str | None, settings: KeySettings, expected: Outcome):
assert await _decide(token, route=route, settings=settings) == (expected, ())
async def test_credential_decisions_do_no_io(
route: str, token: str | None, settings: KeySettings, expected: KeyAuthentication
):
assert await _authenticate(token, route=route, settings=settings) == (expected, ())
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"])
@pytest.mark.parametrize("route", ["/chat/completions", MASTER_ONLY_ROUTE])
async def test_master_key_resolves_end_user_without_key_prefetch(route: str):
handler: Final = RecordingHandler()
assert await _decide(MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",))
assert await _authenticate(MASTER, route=route, handler=handler) == (MasterKey(), ("resolve_end_user",))
assert handler.calls == (("resolve_end_user", False),)
@pytest.mark.asyncio
async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked():
async def test_unknown_virtual_key_is_loaded_from_the_db_then_verified():
handler: Final = RecordingHandler()
outcome, calls = await _decide("sk-x", handler=handler)
assert outcome == _allowed()
assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED)
authentication, calls = await _authenticate("sk-x", handler=handler)
assert authentication == VirtualKey(key=PLAIN_KEY)
assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", "prepare_key", "inspect_key")
assert handler.calls[2] == ("load_key", hash_token("sk-x"))
assert handler.calls[3] == ("prepare_key", True)
@pytest.mark.asyncio
async def test_cached_virtual_key_is_checked_without_a_db_load():
async def test_cached_virtual_key_is_verified_without_a_db_load():
handler: Final = RecordingHandler(presented=NonAdminKey())
outcome, calls = await _decide("sk-x", handler=handler)
assert outcome == _allowed()
assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED)
authentication, calls = await _authenticate("sk-x", handler=handler)
assert authentication == VirtualKey(key=PLAIN_KEY)
assert calls == ("resolve_end_user", "resolve_presented_key", "prepare_key", "inspect_key")
assert handler.calls[2] == ("prepare_key", False)
@pytest.mark.asyncio
@pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW])
async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(expires: datetime | None):
no_db: Final = KeySettings(MASTER, False, False, frozenset({"/global/spend/reset"}), 0.0)
outcome, calls = await _decide(
"sk-admin",
route="/global/spend/reset",
settings=no_db,
handler=RecordingHandler(presented=AdminKey(expires=expires)),
async def test_unexpired_cached_admin_key_skips_the_db_and_key_checks(expires: datetime | None):
authentication, calls = await _authenticate(
"sk-admin", route=MASTER_ONLY_ROUTE, settings=NO_DB, handler=RecordingHandler(presented=AdminKey(expires))
)
assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key"))
assert (authentication, calls) == (CachedAdminKey(), ("resolve_end_user", "resolve_presented_key"))
@pytest.mark.asyncio
async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key():
expires: Final = NOW - timedelta(seconds=1)
handler: Final = RecordingHandler(presented=AdminKey(expires=expires))
outcome, calls = await _decide("sk-admin", handler=handler)
assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin"))
authentication, calls = await _authenticate("sk-admin", handler=handler)
assert authentication == ExpiredKey(expires=expires, now=NOW, key_label="sk-admin")
assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key")
assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin"))
@pytest.mark.asyncio
@pytest.mark.parametrize("presented", [None, NonAdminKey()])
async def test_non_admin_key_cannot_use_master_key_only_routes(presented: PresentedKey | None):
outcome, calls = await _decide("sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented))
assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key"))
@pytest.mark.asyncio
@pytest.mark.parametrize("presented", [None, NonAdminKey()])
async def test_virtual_key_needs_a_database(presented: PresentedKey | None):
no_db: Final = KeySettings(MASTER, False, False, frozenset(), 0.0)
outcome, _ = await _decide("sk-x", settings=no_db, handler=RecordingHandler(presented=presented))
assert outcome == Deny(NoDatabase())
authentication, _ = await _authenticate("sk-x", settings=NO_DB, handler=RecordingHandler(presented=presented))
assert authentication == NoDatabase()
@pytest.mark.asyncio
@pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"])
async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(token: str):
assert await _decide(token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key"))
assert await _authenticate(token) == (NotAVirtualKey(), ("resolve_end_user", "resolve_presented_key"))
@pytest.mark.asyncio
async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with():
expires: Final = NOW - timedelta(seconds=1)
outcome, _ = await _decide("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires)))
assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x")))
authentication, _ = await _authenticate("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires)))
assert authentication == ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x"))
@pytest.mark.asyncio
async def test_blocked_key_is_denied_before_model_access():
assert await _check(replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",))
@pytest.mark.parametrize(
("key", "owner", "expected", "calls"),
[
(replace(PLAIN_KEY, blocked=True), None, KeyBlocked(), ("inspect_key",)),
(
replace(PLAIN_KEY, user_id="u1"),
OwnerFacts(deactivated=True),
OwnerDeactivated(user_id="u1"),
("inspect_key", "load_key_owner"),
),
(replace(PLAIN_KEY, has_token=False), None, KeyWithoutToken(), ("inspect_key",)),
(replace(PLAIN_KEY, user_id="u1"), OwnerFacts(deactivated=False), None, ("inspect_key", "load_key_owner")),
(replace(PLAIN_KEY, user_id="u1"), None, None, ("inspect_key", "load_key_owner")),
],
)
async def test_verify_decides_whether_the_resolved_key_is_usable(
key: KeyFacts, owner: OwnerFacts | None, expected: KeyAuthentication | None, calls: tuple[str, ...]
):
handler: Final = RecordingHandler(key=key, owner=owner)
verified: Final = await PythonCore().keys().verify(handler, _facts("hash"), NOW)
assert verified == (VirtualKey(key=key) if expected is None else expected)
assert _calls(handler) == calls
@pytest.mark.asyncio
async def test_deactivated_owner_is_denied_after_model_access():
outcome, calls = await _check(replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True))
assert (outcome, calls) == (
Deny(OwnerDeactivated(user_id="u1")),
("inspect_key", "enforce_model_access", "load_key_owner"),
async def test_verify_labels_an_expired_key_with_the_hash_it_was_given():
expired: Final = replace(PLAIN_KEY, expires=NOW - timedelta(seconds=1))
verified: Final = await PythonCore().keys().verify(RecordingHandler(key=expired), _facts("hash"), NOW)
assert verified == ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash")
@pytest.mark.asyncio
async def test_expired_key_over_its_team_member_budget_is_denied_as_expired_without_a_budget_alert():
key: Final = replace(
PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"), expires=NOW - timedelta(seconds=1)
)
handler: Final = RecordingHandler(key=key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
authentication, calls = await _authenticate("sk-x", handler=handler)
assert authentication == ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label=hash_token("sk-x"))
assert AUTHORIZE_CALLS.isdisjoint(calls)
@pytest.mark.asyncio
@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)])
async def test_active_or_missing_owner_continues(owner: OwnerFacts | None):
outcome, _ = await _check(replace(PLAIN_KEY, user_id="u1"), owner=owner)
assert outcome == _allowed()
async def test_deactivated_owner_is_denied_before_model_access_is_checked():
handler: Final = RecordingHandler(key=replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True))
authentication, calls = await _authenticate("sk-x", handler=handler)
assert authentication == OwnerDeactivated(user_id="u1")
assert "enforce_model_access" not in calls
@pytest.mark.asyncio
async def test_exempt_request_skips_every_budget_check_but_not_expiry():
@pytest.mark.parametrize(
("handler", "settings", "expected"),
[
(RecordingHandler(), NO_DB, NoDatabase()),
(RecordingHandler(key=replace(PLAIN_KEY, blocked=True)), KEY_SETTINGS, KeyBlocked()),
],
)
async def test_master_key_only_route_no_longer_hides_why_a_key_is_unusable(
handler: RecordingHandler, settings: KeySettings, expected: KeyAuthentication
):
authentication, _ = await _authenticate("sk-x", route=MASTER_ONLY_ROUTE, settings=settings, handler=handler)
assert authentication == expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
("identity", "allow_user_auth", "expected"),
[
(UserAuthRoute(), True, _allowed()),
(UserAuthRoute(), False, Deny(UserAuthDisabled())),
(NoMasterKey(token="sk-x"), False, _allowed()),
(MasterKey(), False, _allowed()),
(CachedAdminKey(), False, _allowed()),
],
)
async def test_non_virtual_key_identities_are_authorized_without_io(
identity: KeyIdentity, allow_user_auth: bool, expected: Outcome
):
settings: Final = replace(AUTHORIZE_SETTINGS, allow_user_auth=allow_user_auth)
assert await _authorize(identity, route=MASTER_ONLY_ROUTE, settings=settings) == (expected, ())
@pytest.mark.asyncio
async def test_virtual_key_on_a_master_key_only_route_is_denied_before_any_io():
assert await _authorize(VirtualKey(key=PLAIN_KEY), route=MASTER_ONLY_ROUTE) == (Deny(MasterKeyOnlyRoute()), ())
@pytest.mark.asyncio
async def test_plain_virtual_key_runs_every_authorization_step_in_order():
assert await _authorize_key(PLAIN_KEY) == (_allowed(), AUTHORIZED)
@pytest.mark.asyncio
async def test_exempt_request_skips_every_budget_check():
key: Final = replace(
PLAIN_KEY,
user_id="u1",
@ -359,13 +448,10 @@ async def test_exempt_request_skips_every_budget_check_but_not_expiry():
spend=5.0,
soft_budget=1.0,
)
assert await _check(key, exempt=True) == (
assert await _authorize_key(key, exempt=True) == (
_allowed(),
("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption", "load_key_context"),
("enforce_model_access", "check_budget_exemption", "load_key_context"),
)
expired: Final = replace(key, expires=NOW - timedelta(seconds=1))
outcome, _ = await _check(expired, exempt=True)
assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash"))
@pytest.mark.asyncio
@ -380,7 +466,7 @@ async def test_exempt_request_skips_every_budget_check_but_not_expiry():
)
async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMemberBudget | None, denied: bool):
key: Final = replace(PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"))
outcome, calls = await _check(key, team_member_budget=budget)
outcome, calls = await _authorize_key(key, team_member_budget=budget)
if denied:
assert budget is not None
assert outcome == Deny(
@ -394,26 +480,20 @@ async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMe
@pytest.mark.asyncio
async def test_team_member_budget_is_read_only_for_a_team_member():
key: Final = replace(PLAIN_KEY, user_id="u1")
_, calls = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
_, calls = await _authorize_key(
replace(PLAIN_KEY, user_id="u1"), team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)
)
assert "read_team_member_budget" not in calls
@pytest.mark.asyncio
async def test_team_member_budget_needs_a_database():
key: Final = replace(PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"))
no_db: Final = KeySettings(MASTER, False, False, frozenset(), 0.0)
_, calls = await _check(key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
assert "read_team_member_budget" not in calls
@pytest.mark.asyncio
async def test_team_member_budget_is_checked_before_expiry():
key: Final = replace(
PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"), expires=NOW - timedelta(seconds=1)
handler: Final = RecordingHandler(key=key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
_, calls = await _authorize(
VirtualKey(key=key), settings=replace(AUTHORIZE_SETTINGS, has_db=False), handler=handler
)
outcome, _ = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5))
assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5))
assert "read_team_member_budget" not in calls
@pytest.mark.asyncio
@ -425,18 +505,8 @@ async def test_team_member_budget_is_checked_before_expiry():
(4.9, 5.0, None, _allowed()),
(6.0, 5.0, 0.5, _allowed(budget_throttle_pct=0.5)),
(4.9, 5.0, 0.5, _allowed()),
(
1e9,
math.inf,
None,
_allowed(),
),
(
1e9,
math.nan,
None,
_allowed(),
),
(1e9, math.inf, None, _allowed()),
(1e9, math.nan, None, _allowed()),
],
)
async def test_key_budget_denies_or_throttles_at_the_limit(
@ -444,17 +514,19 @@ async def test_key_budget_denies_or_throttles_at_the_limit(
):
key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle)
handler: Final = RecordingHandler(key=key, key_spend=spend)
outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), SETTINGS, NOW), handler)
outcome, calls = await _authorize(VirtualKey(key=key), handler=handler)
assert outcome == expected
assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget")
assert handler.calls[5] == ("alert_key_budget", spend)
assert calls[2:5] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget")
assert handler.calls[4] == ("alert_key_budget", spend)
assert ("read_model_budget_scope" in calls) == isinstance(expected, Allow)
@pytest.mark.asyncio
async def test_key_budget_is_enforced_only_on_llm_routes():
key: Final = replace(PLAIN_KEY, max_budget=1.0)
outcome, calls = await _check(key, is_llm_api_route=False, key_spend=2.0)
outcome, calls = await _authorize(
VirtualKey(key=key), is_llm_api_route=False, handler=RecordingHandler(key=key, key_spend=2.0)
)
assert outcome == _allowed()
assert "read_key_spend" not in calls
assert "read_model_budget_scope" in calls
@ -475,16 +547,16 @@ async def test_key_budget_is_enforced_only_on_llm_routes():
async def test_soft_budget_alerts_at_the_threshold_before_model_budgets(
spend: float, soft_budget: float | None, alerted: bool
):
outcome, calls = await _check(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget))
outcome, calls = await _authorize_key(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget))
assert outcome == _allowed()
soft_alert: Final = ("alert_key_soft_budget",) if alerted else ()
assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context")
assert calls[2:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context")
@pytest.mark.asyncio
async def test_throttled_key_still_gets_soft_and_model_budget_checks():
key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0)
outcome, calls = await _check(key, key_spend=6.0)
outcome, calls = await _authorize_key(key, key_spend=6.0)
assert outcome == _allowed(budget_throttle_pct=0.5)
assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "read_model_budget_scope", "load_key_context")
@ -510,31 +582,20 @@ async def test_throttled_key_still_gets_soft_and_model_budget_checks():
async def test_proxy_budget_only_alerts_and_runs_for_exempt_requests_too(
exempt: bool, proxy_max_budget: float, has_db: bool, proxy_spend: float | None, tail: tuple[Call, ...]
):
settings: Final = replace(SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db)
settings: Final = replace(AUTHORIZE_SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db)
handler: Final = RecordingHandler(exempt=exempt, proxy_spend=proxy_spend)
outcome, _ = await _run(PythonCore().keys().check(handler, _facts("hash"), settings, NOW), handler)
outcome, _ = await _authorize(VirtualKey(key=PLAIN_KEY), settings=settings, handler=handler)
assert outcome == _allowed()
assert handler.calls[-len(tail) :] == tail
@pytest.mark.asyncio
@pytest.mark.parametrize("exempt", [False, True])
async def test_key_without_token_is_denied_after_its_context_loads(exempt: bool):
settings: Final = replace(SETTINGS, proxy_max_budget=10.0)
handler: Final = RecordingHandler(key=replace(PLAIN_KEY, has_token=False), exempt=exempt, proxy_spend=1.0)
outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), settings, NOW), handler)
assert outcome == Deny(KeyWithoutToken())
assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget")
@pytest.mark.asyncio
async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_loads():
async def test_a_model_budget_denial_stops_authorization_before_the_key_context_loads():
exceeded: Final = ModelBudgetExceeded(
message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u-1"
)
scope: Final = replace(NO_MODEL_BUDGETS, models=("gpt-4o",), user_has_model_budgets=True)
handler: Final = RecordingHandler(model_budget_scope=scope, user_model_budget_exceeded=exceeded)
outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), SETTINGS, NOW), handler)
outcome, calls = await _authorize_key(PLAIN_KEY, model_budget_scope=scope, user_model_budget_exceeded=exceeded)
assert outcome == Deny(exceeded)
assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget")