refactor(auth): run the common checks through authorize().check_request()

run_centralized_common_checks keeps its signature for its six callers but now
builds a host and asks the core. check_request() decides which requests are
exempt and the order of every step: end-user resolution, the team, user,
project and end-user lookups, the proxy admin user override, the end-user
default budget, common_checks, the team, team member, end-user and sub-call
key budgets, committed spend and the budget reservation, releasing the spend
counter batch even when a check refuses the request.

The public-route exemption now uses the planner's is_public_route instead of
its own copy, so it follows require_auth_for_metrics_endpoint like the planner
does. common_checks and the per-model budget checks still run inside host
methods that raise the proxy's errors. gateway-auth-core mirrors the facet,
host and types.
This commit is contained in:
Yujong Lee 2026-11-10 23:00:37 +09:00
parent 9a79b2b2a9
commit 1edf8b268c
17 changed files with 1001 additions and 401 deletions

View file

@ -81,3 +81,15 @@ pub struct Introspection {
pub role: Value,
pub team_id: Value,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EndUserScope {
pub has_end_user: bool,
pub has_key_end_user_budget: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TokenRole {
pub is_proxy_admin: bool,
pub is_builtin_admin: bool,
}

View file

@ -3,7 +3,7 @@ use std::{future::Future, time::SystemTime};
use crate::{
host::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
OAuth2ProxyHost, RequestCheckHost, TeamMemberModelBudgetHost,
},
identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication,
@ -11,10 +11,10 @@ use crate::{
OAuth2ProxyIdentity, VirtualKey,
},
inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings,
OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
AuthorizeSettings, CheckFacts, CheckSettings, JwtSettings, KeySettings,
OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
},
outcome::{ModelBudgetResult, Outcome},
outcome::{ModelBudgetResult, Outcome, RequestCheck},
validators::{CredentialPlan, KeyPlan, Plan},
};
@ -105,6 +105,12 @@ pub trait AuthorizeCore: Sync {
facts: &RequestFacts,
settings: &AuthorizeSettings,
) -> impl Future<Output = Result<Outcome, H::Error>> + Send;
fn check_request<H: RequestCheckHost>(
&self,
host: &H,
facts: &CheckFacts,
settings: &CheckSettings,
) -> impl Future<Output = Result<RequestCheck, H::Error>> + Send;
}
pub trait AuthCore {

View file

@ -2,8 +2,9 @@ use std::{error::Error, future::Future};
use crate::{
answers::{
FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OwnerFacts, PresentedKey,
SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
EndUserScope, FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OwnerFacts,
PresentedKey, SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
TokenRole,
},
identity::{CustomAuthResult, JwtIdentity},
inputs::RequestFacts,
@ -152,3 +153,48 @@ pub trait CustomAuthHost: Sync {
facts: &RequestFacts,
) -> impl Future<Output = Result<CustomAuthResult, Self::Error>> + Send;
}
pub trait RequestCheckHost: Sync {
type Error: Error + Send + Sync + 'static;
fn resolve_end_user(
&self,
resolve: bool,
) -> impl Future<Output = Result<EndUserScope, Self::Error>> + Send;
fn load_check_objects(&self) -> impl Future<Output = Result<TokenRole, Self::Error>> + Send;
fn override_admin_user(&self, keep_budget_row: bool);
fn apply_project(&self);
fn apply_end_user_default_budget(&self)
-> impl Future<Output = Result<(), Self::Error>> + Send;
fn load_end_user_model_budget(
&self,
assign: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn skips_budget_checks(&self) -> bool;
fn load_team_member_default_model_budget(
&self,
skip_budget_checks: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn merge_request_tags(&self);
fn bind_admission_counters(&self);
fn run_common_checks(
&self,
skip_budget_checks: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn check_team_model_budget(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn check_team_member_model_budget(
&self,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn check_end_user_model_budget(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn check_key_max_budget(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn committed_spend_tracked(&self) -> bool;
fn carry_committed_spend(
&self,
read_committed: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn reserve_budget(
&self,
skip_budget_checks: bool,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn release_spend_counter_batch(&self);
}

View file

@ -73,3 +73,21 @@ pub struct AuthorizeSettings {
pub master_key_only_routes: BTreeSet<String>,
pub proxy_max_budget: f64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CheckFacts {
pub request: RequestFacts,
pub force_virtual_key_checks: bool,
pub sub_call: bool,
pub token_has_end_user: bool,
pub route_skips_budget_checks: bool,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CheckSettings {
pub plan: PlanSettings,
pub auth_off_pass_through_routes: BTreeSet<String>,
pub no_auth_dev_mode: bool,
pub custom_auth_skips_checks: bool,
pub has_db: bool,
}

View file

@ -8,8 +8,9 @@ mod outcome;
mod validators;
pub use answers::{
FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope, OwnerFacts,
PresentedKey, SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
EndUserScope, FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope,
OwnerFacts, PresentedKey, SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope,
TokenRole,
};
pub use facets::{
AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore,
@ -17,7 +18,7 @@ pub use facets::{
};
pub use host::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
OAuth2ProxyHost, RequestCheckHost, TeamMemberModelBudgetHost,
};
pub use identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, CustomIdentity, Identity,
@ -25,9 +26,13 @@ pub use identity::{
OAuth2ProxyDenial, OAuth2ProxyIdentity, VirtualKey,
};
pub use inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings,
PlanFacts, PlanSettings, RequestFacts, RoutingOverride,
AuthorizeSettings, CheckFacts, CheckSettings, JwtSettings, KeySettings,
OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
RoutingOverride,
};
pub use native::{NativeCore, Unported};
pub use outcome::{Denial, ModelBudgetExceeded, ModelBudgetResult, Outcome, WithinModelBudgets};
pub use outcome::{
Denial, ExemptionReason, ModelBudgetExceeded, ModelBudgetResult, Outcome, RequestCheck,
WithinModelBudgets,
};
pub use validators::{CredentialPlan, KeyPlan, OAuth2Mode, Plan, PremiumFeature};

View file

@ -7,7 +7,7 @@ use crate::{
},
host::{
AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost,
OAuth2ProxyHost, TeamMemberModelBudgetHost,
OAuth2ProxyHost, RequestCheckHost, TeamMemberModelBudgetHost,
},
identity::{
Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication,
@ -15,10 +15,10 @@ use crate::{
OAuth2ProxyIdentity, VirtualKey,
},
inputs::{
AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings,
OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
AuthorizeSettings, CheckFacts, CheckSettings, JwtSettings, KeySettings,
OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts,
},
outcome::{ModelBudgetResult, Outcome},
outcome::{ModelBudgetResult, Outcome, RequestCheck},
validators::{CredentialPlan, KeyPlan, Plan},
};
@ -202,4 +202,13 @@ impl AuthorizeCore for Unported {
) -> Result<Outcome, H::Error> {
unimplemented!("{RUNS_IN_PYTHON}")
}
async fn check_request<H: RequestCheckHost>(
&self,
_host: &H,
_facts: &CheckFacts,
_settings: &CheckSettings,
) -> Result<RequestCheck, H::Error> {
unimplemented!("{RUNS_IN_PYTHON}")
}
}

View file

@ -34,3 +34,17 @@ pub enum Outcome {
Allow { budget_throttle_pct: Option<f64> },
Deny(Denial),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)]
pub enum ExemptionReason {
PublicRoute,
PassThroughAuthOff,
NoAuthDevMode,
CustomAuth,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RequestCheck {
Exempt(ExemptionReason),
Checked,
}

View file

@ -12,7 +12,7 @@ Modules here import only the standard library and each other. A core does no IO
`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()` 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.
`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().authorize()` checks a key identity, and `authorize().check_request()` then runs the common checks every authenticated request passes, whatever authenticated it. `check_request()` decides which requests are exempt (public routes by the planner's own rule, pass-through routes without `auth: true`, no-auth dev mode, custom auth that opts out) and the order of every step, while `common_checks` itself and the per-model budget checks still run inside host methods that raise the proxy's errors. Moving those decisions into the core is the next step.
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`.

View file

@ -88,3 +88,15 @@ class Introspection:
user_id: object
role: object
team_id: object
@dataclass(frozen=True, slots=True)
class EndUserScope:
has_end_user: bool
has_key_end_user_budget: bool
@dataclass(frozen=True, slots=True)
class TokenRole:
is_proxy_admin: bool
is_builtin_admin: bool

View file

@ -9,6 +9,7 @@ from litellm.proxy.auth.core.host import (
ModelBudgetHost,
OAuth2IntrospectionHost,
OAuth2ProxyHost,
RequestCheckHost,
TeamMemberModelBudgetHost,
)
from litellm.proxy.auth.core.identity import (
@ -28,6 +29,8 @@ from litellm.proxy.auth.core.identity import (
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
CheckFacts,
CheckSettings,
JwtSettings,
KeySettings,
OAuth2IntrospectionSettings,
@ -36,7 +39,7 @@ from litellm.proxy.auth.core.inputs import (
PlanSettings,
RequestFacts,
)
from litellm.proxy.auth.core.outcome import ModelBudgetResult, Outcome
from litellm.proxy.auth.core.outcome import ModelBudgetResult, Outcome, RequestCheck
from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan
@ -95,6 +98,10 @@ class AuthorizeCore(Protocol):
self, host: AuthorizeHost, identity: KeyIdentity, facts: RequestFacts, settings: AuthorizeSettings
) -> Outcome: ...
async def check_request(
self, host: RequestCheckHost, facts: CheckFacts, settings: CheckSettings
) -> RequestCheck: ...
class AuthCore(Protocol):
def keys(self) -> KeyCore: ...

View file

@ -1,6 +1,7 @@
from typing import Protocol
from litellm.proxy.auth.core.answers import (
EndUserScope,
FallbackTarget,
Introspection,
KeyFacts,
@ -11,6 +12,7 @@ from litellm.proxy.auth.core.answers import (
TeamMember,
TeamMemberBudget,
TeamMemberModelBudgetScope,
TokenRole,
)
from litellm.proxy.auth.core.identity import CustomAuthResult, JwtIdentity
from litellm.proxy.auth.core.inputs import RequestFacts
@ -101,3 +103,43 @@ class OAuth2ProxyHost(Protocol):
class CustomAuthHost(Protocol):
async def custom_auth(self, facts: RequestFacts) -> CustomAuthResult: ...
class RequestCheckHost(Protocol):
async def resolve_end_user(self, resolve: bool) -> EndUserScope: ...
async def load_check_objects(self) -> TokenRole: ...
def override_admin_user(self, keep_budget_row: bool) -> None: ...
def apply_project(self) -> None: ...
async def apply_end_user_default_budget(self) -> None: ...
async def load_end_user_model_budget(self, assign: bool) -> None: ...
def skips_budget_checks(self) -> bool: ...
async def load_team_member_default_model_budget(self, skip_budget_checks: bool) -> None: ...
def merge_request_tags(self) -> None: ...
def bind_admission_counters(self) -> None: ...
async def run_common_checks(self, skip_budget_checks: bool) -> None: ...
async def check_team_model_budget(self) -> None: ...
async def check_team_member_model_budget(self) -> None: ...
async def check_end_user_model_budget(self) -> None: ...
async def check_key_max_budget(self) -> None: ...
def committed_spend_tracked(self) -> bool: ...
async def carry_committed_spend(self, read_committed: bool) -> None: ...
async def reserve_budget(self, skip_budget_checks: bool) -> None: ...
def release_spend_counter_batch(self) -> None: ...

View file

@ -73,3 +73,21 @@ class AuthorizeSettings:
has_db: bool
master_key_only_routes: frozenset[str]
proxy_max_budget: float
@dataclass(frozen=True, slots=True)
class CheckFacts:
request: RequestFacts
force_virtual_key_checks: bool
sub_call: bool
token_has_end_user: bool
route_skips_budget_checks: bool
@dataclass(frozen=True, slots=True)
class CheckSettings:
plan: PlanSettings
auth_off_pass_through_routes: frozenset[str]
no_auth_dev_mode: bool
custom_auth_skips_checks: bool
has_db: bool

View file

@ -1,4 +1,5 @@
from dataclasses import dataclass
from enum import Enum
from typing import TypeAlias
@ -58,3 +59,23 @@ class Deny:
Outcome: TypeAlias = Allow | Deny
class ExemptionReason(Enum):
PUBLIC_ROUTE = "public_route"
PASS_THROUGH_AUTH_OFF = "pass_through_auth_off"
NO_AUTH_DEV_MODE = "no_auth_dev_mode"
CUSTOM_AUTH = "custom_auth"
@dataclass(frozen=True, slots=True)
class Exempt:
reason: ExemptionReason
@dataclass(frozen=True, slots=True)
class Checked:
pass
RequestCheck: TypeAlias = Exempt | Checked

View file

@ -5,6 +5,7 @@ from typing import Final
from litellm.proxy.auth.core import oauth2 as oauth2_rules
from litellm.proxy.auth.core import plan as plan_rules
from litellm.proxy.auth.core import request_checks
from litellm.proxy.auth.core.answers import AdminKey, KeyFacts, TeamMember
from litellm.proxy.auth.core.core import (
AnonymousCore,
@ -24,6 +25,7 @@ from litellm.proxy.auth.core.host import (
ModelBudgetHost,
OAuth2IntrospectionHost,
OAuth2ProxyHost,
RequestCheckHost,
TeamMemberModelBudgetHost,
)
from litellm.proxy.auth.core.identity import (
@ -53,6 +55,8 @@ from litellm.proxy.auth.core.identity import (
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
CheckFacts,
CheckSettings,
KeySettings,
OAuth2IntrospectionSettings,
OAuth2ProxySettings,
@ -74,6 +78,7 @@ from litellm.proxy.auth.core.outcome import (
ModelBudgetExceeded,
ModelBudgetResult,
Outcome,
RequestCheck,
TeamMemberBudgetExceeded,
UserAuthDisabled,
WithinModelBudgets,
@ -184,6 +189,9 @@ class PythonAuthorizeCore(AuthorizeCore):
case VirtualKey(key=key):
return await _authorize_virtual_key(host, key, facts, settings)
async def check_request(self, host: RequestCheckHost, facts: CheckFacts, settings: CheckSettings) -> RequestCheck:
return await request_checks.check_request(host, facts, settings)
class PythonModelBudgetCore(ModelBudgetCore):
async def check_key_and_user(self, host: ModelBudgetHost) -> ModelBudgetResult:

View file

@ -0,0 +1,60 @@
from typing import Final
from litellm.proxy.auth.core.host import RequestCheckHost
from litellm.proxy.auth.core.inputs import CheckFacts, CheckSettings
from litellm.proxy.auth.core.outcome import Checked, Exempt, ExemptionReason, RequestCheck
from litellm.proxy.auth.core.plan import is_public_route
def _exemption(facts: CheckFacts, settings: CheckSettings) -> ExemptionReason | None:
route: Final = facts.request.route
if facts.force_virtual_key_checks:
return None
if is_public_route(route, settings.plan):
return ExemptionReason.PUBLIC_ROUTE
if route in settings.auth_off_pass_through_routes:
return ExemptionReason.PASS_THROUGH_AUTH_OFF
if settings.no_auth_dev_mode:
return ExemptionReason.NO_AUTH_DEV_MODE
if settings.custom_auth_skips_checks:
return ExemptionReason.CUSTOM_AUTH
return None
async def check_request(host: RequestCheckHost, facts: CheckFacts, settings: CheckSettings) -> RequestCheck:
exemption: Final = _exemption(facts, settings)
if exemption is not None:
return Exempt(exemption)
end_user: Final = await host.resolve_end_user(resolve=not facts.token_has_end_user)
role: Final = await host.load_check_objects()
if role.is_proxy_admin:
host.override_admin_user(keep_budget_row=not role.is_builtin_admin)
host.apply_project()
if end_user.has_end_user and end_user.has_key_end_user_budget and settings.has_db:
await host.apply_end_user_default_budget()
await host.load_end_user_model_budget(assign=end_user.has_end_user)
skip_budget_checks: Final = host.skips_budget_checks()
await host.load_team_member_default_model_budget(
skip_budget_checks=skip_budget_checks or facts.route_skips_budget_checks
)
host.merge_request_tags()
host.bind_admission_counters()
try:
await host.run_common_checks(skip_budget_checks=skip_budget_checks)
if not skip_budget_checks:
await _check_model_budgets(host, facts)
if host.committed_spend_tracked():
await host.carry_committed_spend(read_committed=not facts.route_skips_budget_checks)
await host.reserve_budget(skip_budget_checks=skip_budget_checks or facts.sub_call)
finally:
host.release_spend_counter_batch()
return Checked()
async def _check_model_budgets(host: RequestCheckHost, facts: CheckFacts) -> None:
await host.check_team_model_budget()
if not facts.route_skips_budget_checks:
await host.check_team_member_model_budget()
await host.check_end_user_model_budget()
if facts.sub_call and facts.request.is_llm_api_route:
await host.check_key_max_budget()

View file

@ -121,11 +121,11 @@ from litellm.proxy.auth.auth_utils import (
pre_db_read_auth_checks,
request_dispatched_to_pass_through_endpoint,
request_dispatched_to_provider_pass_through,
route_in_additonal_public_routes,
)
from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded
from litellm.proxy.auth.core.answers import (
AdminKey,
EndUserScope,
FallbackTarget,
KeyFacts,
KeyModelBudgets,
@ -136,6 +136,7 @@ from litellm.proxy.auth.core.answers import (
TeamMember,
TeamMemberBudget,
TeamMemberModelBudgetScope,
TokenRole,
)
from litellm.proxy.auth.core.credential import (
PresentedKeys,
@ -164,6 +165,8 @@ from litellm.proxy.auth.core.identity import (
)
from litellm.proxy.auth.core.inputs import (
AuthorizeSettings,
CheckFacts,
CheckSettings,
KeySettings,
PlanFacts,
PlanSettings,
@ -3775,11 +3778,13 @@ async def run_centralized_common_checks(
"""
from litellm.proxy.proxy_server import (
general_settings,
jwt_handler,
litellm_proxy_admin_name,
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
master_key,
model_max_budget_limiter,
premium_user,
prisma_client,
proxy_logging_obj,
spend_counter_store,
@ -3787,420 +3792,462 @@ async def run_centralized_common_checks(
user_custom_auth,
)
# Public routes (e.g. /health/liveness) are exempt from
# auth in the builder — the wrapper must not retroactively apply
# authz on top, or k8s readiness probes and other unauthenticated
# callers get 401.
if not force_virtual_key_checks and (
route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route)
):
return
# User-configured pass-through endpoints with ``auth: false`` are
# explicitly unauthenticated — the builder returns an empty
# UserAPIKeyAuth() and the request is forwarded as-is. Running
# common_checks on the empty token would reject the request as
# admin-only. The "auth" flag on the endpoint config is the
# contract; honor it.
pass_through_endpoints: Final = general_settings.get("pass_through_endpoints", None)
if not force_virtual_key_checks and pass_through_endpoints is not None:
for endpoint in pass_through_endpoints:
if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True:
return
# No-auth dev mode: master_key unset AND no JWT/OAuth2 auth
# configured. The builder returns an INTERNAL_USER token for any
# api_key; the proxy is unauthenticated by configuration.
# Running common_checks would block every admin route on these
# deployments where that was previously not the contract. If any
# authn is enabled (JWT, OAuth2, OAuth2-proxy), authz must run.
if not force_virtual_key_checks and is_no_auth_dev_mode(master_key, general_settings):
return
if (
not force_virtual_key_checks
and user_custom_auth is not None
and not general_settings.get("custom_auth_run_common_checks", False)
):
return
parent_otel_span: Final = user_api_key_auth_obj.parent_otel_span
# In the integrated auth flow ``_user_api_key_auth_builder`` has already
# resolved the end-user id and attached it here. Reuse that to avoid a
# second extraction pass; fall back to extracting locally when the
# function is invoked in isolation (e.g. in direct unit tests).
key_end_user_budget_id: Final = get_key_end_user_budget_id(user_api_key_auth_obj.metadata)
end_user_id = user_api_key_auth_obj.end_user_id
if end_user_id is None:
raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, safe_get_request_headers(request))
end_user_id = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
settings: Final = CheckSettings(
plan=_plan_settings(
general_settings, jwt_handler, premium_user=premium_user is True, custom_auth=user_custom_auth is not None
),
auth_off_pass_through_routes=frozenset(
str(endpoint.get("path", ""))
for endpoint in pass_through_endpoints or ()
if isinstance(endpoint, dict) and endpoint.get("auth") is not True
),
no_auth_dev_mode=is_no_auth_dev_mode(master_key, general_settings),
custom_auth_skips_checks=user_custom_auth is not None
and not general_settings.get("custom_auth_run_common_checks", False),
has_db=prisma_client is not None,
)
facts: Final = CheckFacts(
request=RequestFacts(
route=route,
key_end_user_budget_id=key_end_user_budget_id,
)
if end_user_id is not None and key_end_user_budget_id is not None:
user_api_key_auth_obj.end_user_id = end_user_id
fetch_coros: Final = []
if user_api_key_auth_obj.team_id is not None and user_api_key_auth_obj.team_id != UI_TEAM_ID:
fetch_coros.append(
_safe_fetch(
"team",
get_team_object(
team_id=user_api_key_auth_obj.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("team", _noop_none()))
if user_api_key_auth_obj.user_id is not None:
fetch_coros.append(
_safe_fetch(
"user",
get_user_object(
user_id=user_api_key_auth_obj.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("user", _noop_none()))
if user_api_key_auth_obj.project_id is not None:
fetch_coros.append(
_safe_fetch(
"project",
get_project_object(
project_id=user_api_key_auth_obj.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("project", _noop_none()))
if end_user_id:
fetch_coros.append(
_safe_fetch(
"end_user",
get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
token_end_user_max_budget=user_api_key_auth_obj.end_user_max_budget,
key_end_user_budget_id=key_end_user_budget_id,
),
)
)
else:
fetch_coros.append(_safe_fetch("end_user", _noop_none()))
fetch_coros.append(
_safe_fetch(
"global_spend",
get_global_proxy_spend(
litellm_proxy_admin_name=litellm_proxy_admin_name,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
token=user_api_key_auth_obj.token or "",
proxy_logging_obj=proxy_logging_obj,
),
)
token=user_api_key_auth_obj.token,
is_llm_api_route=RouteChecks.is_llm_api_route(route=route),
),
force_virtual_key_checks=force_virtual_key_checks,
sub_call=sub_call,
token_has_end_user=user_api_key_auth_obj.end_user_id is not None,
route_skips_budget_checks=route_skips_budget_checks(route=route),
)
# Per-fetch error isolation. ``_safe_fetch`` lets HTTPException,
# ProxyException, and BudgetExceededError escape (everything else is
# already swallowed to None). A bare ``except`` over ``gather`` would
# let one fetch's HTTPException null out every other context — e.g.
# a 404 from ``get_team_object`` (token references a deleted team)
# would silently skip the user, end-user, project, and global-spend
# checks. Use ``return_exceptions=True`` and apply per-fetch fallback
# so a missing team only zeros out the team object.
(
team_result,
user_result,
project_result,
end_user_result,
global_spend_result,
) = await asyncio.gather(*fetch_coros, return_exceptions=True)
# ProxyException / BudgetExceededError are authorization failures —
# propagate so the wrapper renders them. HTTPException is fallback
# material (404 from get_team_object is the only known producer).
for r in (
team_result,
user_result,
project_result,
end_user_result,
global_spend_result,
):
if isinstance(r, (ProxyException, litellm.BudgetExceededError)):
raise r
# Use BaseException (not HTTPException) in the narrowing checks so
# mypy can narrow ``Any | BaseException`` to the typed object in the
# else branch. After the for-loop above, the only BaseException that
# can still appear here is HTTPException (other listed re-raises were
# propagated; non-listed exceptions were already swallowed to None).
team_object: LiteLLM_TeamTableCachedObj | None
if isinstance(team_result, BaseException):
# Token-derived fallback only valid when a team_id is set;
# _team_obj_from_token asserts that precondition.
if user_api_key_auth_obj.team_id is None:
team_object = None
elif _token_can_vouch_for_team(user_api_key_auth_obj, team_result):
team_object = _team_obj_from_token(user_api_key_auth_obj)
else:
raise team_result
else:
team_object = (
_team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id == UI_TEAM_ID else team_result
)
user_object: LiteLLM_UserTable | None = user_result if isinstance(user_result, LiteLLM_UserTable) else None
project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
None if isinstance(project_result, BaseException) else project_result
)
end_user_object: Final = end_user_result if isinstance(end_user_result, LiteLLM_EndUserTable) else None
global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result
carry_team_and_user_budget_state(
valid_token=user_api_key_auth_obj,
team_object=team_object,
user_object=user_object,
)
await _inherit_org_identity(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=team_object,
host: Final = _RequestCheckHost(
token=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
admin_user_id=litellm_proxy_admin_name,
general_settings=general_settings,
llm_router=llm_router,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
model_max_budget_limiter=model_max_budget_limiter,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
spend_counter_store=spend_counter_store,
user_api_key_cache=user_api_key_cache,
keep_token_limits=user_custom_auth is not None,
)
_ = await _AUTH_CORE.authorize().check_request(host, facts, settings)
# common_checks identifies admin via user_object, not the token
# (non_proxy_admin_allowed_routes_check). JWT admin shortcut and
# master_key tokens get admin from the token; the DB row for the
# same user_id (e.g. litellm_proxy_admin_name = "default_user_id")
# may have a non-admin user_role and would otherwise demote the
# caller. The token is the source of truth for these paths — force
# the admin user_object whenever the token says PROXY_ADMIN, even
# if a DB row was fetched. Every admin token except the built-in admin
# identity (litellm_proxy_admin_name: the master key and the UI
# username/password login) keeps the row's personal budget (max_budget,
# and user_email for budget alerts) so common_checks enforces it; every
# other field keeps its default.
if user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN:
budget_row: Final = (
user_object
if isinstance(user_object, LiteLLM_UserTable) and user_api_key_auth_obj.user_id != litellm_proxy_admin_name
else None
class _RequestCheckHost:
def __init__(
self,
*,
token: UserAPIKeyAuth,
request: Request,
request_data: dict[str, object], # mutable-ok: common checks merge tags and rewrite the model in place
route: str,
admin_user_id: str,
general_settings: dict, # pyright: ignore[reportMissingTypeArgument] # proxy global is an untyped dict
llm_router: litellm.Router | None,
llm_model_list: Sequence[object] | None,
model_max_budget_limiter: _EveryModelBudgetLimiter,
prisma_client: PrismaClient | None,
proxy_logging_obj: ProxyLogging,
spend_counter_store: SpendCounterStore,
user_api_key_cache: UserApiKeyCache,
keep_token_limits: bool,
) -> None:
self._token: Final = token
self._request: Final = request
self._request_data: Final = request_data
self._route: Final = route
self._admin_user_id: Final = admin_user_id
self._general_settings: Final = general_settings
self._llm_router: Final = llm_router
self._llm_model_list: Final = llm_model_list
self._limiter: Final = model_max_budget_limiter
self._prisma_client: Final = prisma_client
self._proxy_logging_obj: Final = proxy_logging_obj
self._spend_counter_store: Final = spend_counter_store
self._user_api_key_cache: Final = user_api_key_cache
self._keep_token_limits: Final = keep_token_limits
self._parent_otel_span: Final = token.parent_otel_span
self._key_end_user_budget_id: Final = get_key_end_user_budget_id(token.metadata)
self.end_user_id: str | None = token.end_user_id
self.team_object: LiteLLM_TeamTableCachedObj | LiteLLM_TeamTable | None = None
self.user_object: LiteLLM_UserTable | None = None
self.project_object: LiteLLM_ProjectTableCachedObj | None = None
self.end_user_object: LiteLLM_EndUserTable | None = None
self.global_proxy_spend: float | None = None
self.models: list[str] = [] # mutable-ok: the team member check reuses the team check's model list
def _team_table(self) -> LiteLLM_TeamTable | None:
return self.team_object if isinstance(self.team_object, LiteLLM_TeamTable) else None
def _request_models(self) -> list[str]:
return _get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=self._request_data,
route=self._route,
request=self._request,
llm_router=self._llm_router,
team_id=self._token.team_id,
)
)
user_object = LiteLLM_UserTable(
user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name,
async def resolve_end_user(self, resolve: bool) -> EndUserScope:
if resolve:
raw_end_user_id: Final = get_end_user_id_from_request_body(
self._request_data, safe_get_request_headers(self._request)
)
self.end_user_id = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
route=self._route,
key_end_user_budget_id=self._key_end_user_budget_id,
)
if self.end_user_id is not None and self._key_end_user_budget_id is not None:
self._token.end_user_id = self.end_user_id
return EndUserScope(
has_end_user=bool(self.end_user_id), has_key_end_user_budget=self._key_end_user_budget_id is not None
)
async def load_check_objects(self) -> TokenRole:
fetch_coros: Final = []
if self._token.team_id is not None and self._token.team_id != UI_TEAM_ID:
fetch_coros.append(
_safe_fetch(
"team",
get_team_object(
team_id=self._token.team_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("team", _noop_none()))
if self._token.user_id is not None:
fetch_coros.append(
_safe_fetch(
"user",
get_user_object(
user_id=self._token.user_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
user_id_upsert=False,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("user", _noop_none()))
if self._token.project_id is not None:
fetch_coros.append(
_safe_fetch(
"project",
get_project_object(
project_id=self._token.project_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
proxy_logging_obj=self._proxy_logging_obj,
),
)
)
else:
fetch_coros.append(_safe_fetch("project", _noop_none()))
if self.end_user_id:
fetch_coros.append(
_safe_fetch(
"end_user",
get_end_user_object(
end_user_id=self.end_user_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
route=self._route,
token_end_user_max_budget=self._token.end_user_max_budget,
key_end_user_budget_id=self._key_end_user_budget_id,
),
)
)
else:
fetch_coros.append(_safe_fetch("end_user", _noop_none()))
fetch_coros.append(
_safe_fetch(
"global_spend",
get_global_proxy_spend(
litellm_proxy_admin_name=self._admin_user_id,
user_api_key_cache=self._user_api_key_cache,
prisma_client=self._prisma_client,
token=self._token.token or "",
proxy_logging_obj=self._proxy_logging_obj,
),
)
)
# Per-fetch error isolation. ``_safe_fetch`` lets HTTPException,
# ProxyException, and BudgetExceededError escape (everything else is
# already swallowed to None). A bare ``except`` over ``gather`` would
# let one fetch's HTTPException null out every other context — e.g.
# a 404 from ``get_team_object`` (token references a deleted team)
# would silently skip the user, end-user, project, and global-spend
# checks. Use ``return_exceptions=True`` and apply per-fetch fallback
# so a missing team only zeros out the team object.
(
team_result,
user_result,
project_result,
end_user_result,
global_spend_result,
) = await asyncio.gather(*fetch_coros, return_exceptions=True)
# ProxyException / BudgetExceededError are authorization failures —
# propagate so the wrapper renders them. HTTPException is fallback
# material (404 from get_team_object is the only known producer).
for r in (
team_result,
user_result,
project_result,
end_user_result,
global_spend_result,
):
if isinstance(r, (ProxyException, litellm.BudgetExceededError)):
raise r
# Use BaseException (not HTTPException) in the narrowing checks so
# mypy can narrow ``Any | BaseException`` to the typed object in the
# else branch. After the for-loop above, the only BaseException that
# can still appear here is HTTPException (other listed re-raises were
# propagated; non-listed exceptions were already swallowed to None).
team_object: LiteLLM_TeamTableCachedObj | None
if isinstance(team_result, BaseException):
# Token-derived fallback only valid when a team_id is set;
# _team_obj_from_token asserts that precondition.
if self._token.team_id is None:
team_object = None
elif _token_can_vouch_for_team(self._token, team_result):
team_object = _team_obj_from_token(self._token)
else:
raise team_result
else:
team_object = _team_obj_from_token(self._token) if self._token.team_id == UI_TEAM_ID else team_result
user_object: LiteLLM_UserTable | None = user_result if isinstance(user_result, LiteLLM_UserTable) else None
project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
None if isinstance(project_result, BaseException) else project_result
)
end_user_object: Final = end_user_result if isinstance(end_user_result, LiteLLM_EndUserTable) else None
global_proxy_spend: float | None = (
None if isinstance(global_spend_result, BaseException) else global_spend_result
)
carry_team_and_user_budget_state(
valid_token=self._token,
team_object=team_object,
user_object=user_object,
)
await _inherit_org_identity(
user_api_key_auth_obj=self._token,
team_object=team_object,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
)
self.team_object = team_object
self.user_object = user_object
self.project_object = project_object
self.end_user_object = end_user_object
self.global_proxy_spend = global_proxy_spend
return TokenRole(
is_proxy_admin=self._token.user_role == LitellmUserRoles.PROXY_ADMIN,
is_builtin_admin=self._token.user_id == self._admin_user_id,
)
def override_admin_user(self, keep_budget_row: bool) -> None:
user_object: Final = self.user_object
budget_row: Final = user_object if user_object is not None and keep_budget_row else None
self.user_object = LiteLLM_UserTable(
user_id=self._token.user_id or self._admin_user_id,
user_role=LitellmUserRoles.PROXY_ADMIN,
spend=user_object.spend if user_object is not None else 0.0,
max_budget=budget_row.max_budget if budget_row is not None else None,
user_email=budget_row.user_email if budget_row is not None else None,
object_permission_id=(
user_object.object_permission_id if isinstance(user_object, LiteLLM_UserTable) else None
),
object_permission_id=user_object.object_permission_id if user_object is not None else None,
)
if project_object is not None:
user_api_key_auth_obj.project_metadata = project_object.metadata
user_api_key_auth_obj.project_alias = project_object.project_alias
def apply_project(self) -> None:
if self.project_object is not None:
self._token.project_metadata = self.project_object.metadata
self._token.project_alias = self.project_object.project_alias
if end_user_id and key_end_user_budget_id is not None and prisma_client is not None:
async def apply_end_user_default_budget(self) -> None:
key_end_user_budget_id: Final = self._key_end_user_budget_id
prisma_client: Final = self._prisma_client
if key_end_user_budget_id is None or prisma_client is None:
raise RuntimeError("the auth core applied an end-user default budget without a budget id or database")
await _apply_key_end_user_default_budget_to_token(
valid_token=user_api_key_auth_obj,
end_user_object=end_user_object,
valid_token=self._token,
end_user_object=self.end_user_object,
key_end_user_budget_id=key_end_user_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
keep_token_limits=user_custom_auth is not None,
)
end_user_model_max_budget: Final = await _end_user_model_max_budget(
end_user_id=end_user_id,
end_user_object=end_user_object,
key_end_user_budget_id=key_end_user_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
if end_user_id and (
not user_api_key_auth_obj.authenticated_by_custom_auth
or user_api_key_auth_obj.end_user_model_max_budget is None
):
user_api_key_auth_obj.end_user_model_max_budget = ( # rebind-ok: the spend hook charges from this token
dict(end_user_model_max_budget) if end_user_model_max_budget else None
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
keep_token_limits=self._keep_token_limits,
)
skip_budget_checks: Final = _should_skip_budget_checks(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
typed_team_object: Final[LiteLLM_TeamTable | None] = (
team_object if isinstance(team_object, LiteLLM_TeamTable) else None
)
await _load_team_member_default_model_budget(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=typed_team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
skip_budget_checks=skip_budget_checks or route_skips_budget_checks(route=route),
proxy_logging_obj=proxy_logging_obj,
)
async def load_end_user_model_budget(self, assign: bool) -> None:
end_user_model_max_budget: Final = await _end_user_model_max_budget(
end_user_id=self.end_user_id,
end_user_object=self.end_user_object,
key_end_user_budget_id=self._key_end_user_budget_id,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
)
if assign and (not self._token.authenticated_by_custom_auth or self._token.end_user_model_max_budget is None):
self._token.end_user_model_max_budget = (
dict(end_user_model_max_budget) if end_user_model_max_budget else None
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
# any tag merge runs. Without this, header tags from
# apply_client_tag_policy_pre_auth would land in `metadata` while the
# later seed in common_checks pushes key tags and the
# _tag_max_budget_check read into `litellm_metadata`, hiding header
# tags from per-tag budget enforcement on LITELLM_METADATA_ROUTES.
# Merge x-litellm-tags into request_data BEFORE common_checks runs.
# _tag_max_budget_check inside common_checks only inspects request_data;
# without this pre-merge, header-supplied tags bypass tag-budget
# enforcement.
for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_data):
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route)
LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(
request=request,
request_data=tag_target,
user_api_key_dict=user_api_key_auth_obj,
def skips_budget_checks(self) -> bool:
return _should_skip_budget_checks(
request_data=self._request_data,
route=self._route,
request=self._request,
llm_router=self._llm_router,
team_id=self._token.team_id,
)
bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id)
bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id)
try:
_ = await common_checks(
request=request,
request_body=request_data,
team_object=team_object,
user_object=user_object,
end_user_object=end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=user_api_key_auth_obj,
async def load_team_member_default_model_budget(self, skip_budget_checks: bool) -> None:
await _load_team_member_default_model_budget(
user_api_key_auth_obj=self._token,
team_object=self._team_table(),
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
skip_budget_checks=skip_budget_checks,
project_object=project_object,
proxy_logging_obj=self._proxy_logging_obj,
)
if not skip_budget_checks:
models: Final = _get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
)
await _check_team_model_budget(
valid_token=user_api_key_auth_obj,
model_max_budget_limiter=model_max_budget_limiter,
models=models,
)
typed_project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
project_object if isinstance(project_object, LiteLLM_ProjectTableCachedObj) else None
)
if not route_skips_budget_checks(route=route):
await _check_team_member_model_budget_with_fallback(
valid_token=user_api_key_auth_obj,
model_max_budget_limiter=model_max_budget_limiter,
models=models,
request_data=request_data,
request=request,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
project_object=typed_project_object,
team_object=typed_team_object,
end_user_object=end_user_object if isinstance(end_user_object, LiteLLM_EndUserTable) else None,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _check_end_user_model_budget(
end_user_id=end_user_id,
end_user_model_max_budget=user_api_key_auth_obj.end_user_model_max_budget,
model_max_budget_limiter=model_max_budget_limiter,
team_id=user_api_key_auth_obj.team_id,
models=_get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
),
)
if sub_call and RouteChecks.is_llm_api_route(route=route):
await virtual_key_max_budget_check(
valid_token=user_api_key_auth_obj,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_object,
)
if committed_spend_tracked():
await carry_committed_budget_spend(
valid_token=user_api_key_auth_obj,
team_object=typed_team_object,
user_object=user_object,
end_user_object=end_user_object,
read_committed=(
None
if route_skips_budget_checks(route=route)
else partial(read_admission_committed_spend, spend_counter_store)
),
def merge_request_tags(self) -> None:
# Pin the metadata variable name before any tag merge runs, and merge x-litellm-tags into the request
# before common_checks, whose tag budget check only reads request_data.
for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(self._request, self._request_data):
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=self._route)
LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(
request=self._request,
request_data=tag_target,
user_api_key_dict=self._token,
)
def bind_admission_counters(self) -> None:
bind_admission_counter_keys(self._token, end_user_id=self.end_user_id)
bind_committed_admission_counters(self._token, end_user_id=self.end_user_id)
async def run_common_checks(self, skip_budget_checks: bool) -> None:
_ = await common_checks(
request=self._request,
request_body=self._request_data,
team_object=self.team_object,
user_object=self.user_object,
end_user_object=self.end_user_object,
general_settings=self._general_settings,
global_proxy_spend=self.global_proxy_spend,
route=self._route,
llm_router=self._llm_router,
proxy_logging_obj=self._proxy_logging_obj,
valid_token=self._token,
skip_budget_checks=skip_budget_checks,
project_object=self.project_object,
)
async def check_team_model_budget(self) -> None:
self.models = self._request_models()
await _check_team_model_budget(
valid_token=self._token, model_max_budget_limiter=self._limiter, models=self.models
)
async def check_team_member_model_budget(self) -> None:
await _check_team_member_model_budget_with_fallback(
valid_token=self._token,
model_max_budget_limiter=self._limiter,
models=self.models,
request_data=self._request_data,
request=self._request,
llm_model_list=self._llm_model_list,
llm_router=self._llm_router,
project_object=self.project_object,
team_object=self._team_table(),
end_user_object=self.end_user_object,
route=self._route,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
proxy_logging_obj=self._proxy_logging_obj,
)
async def check_end_user_model_budget(self) -> None:
await _check_end_user_model_budget(
end_user_id=self.end_user_id,
end_user_model_max_budget=self._token.end_user_model_max_budget,
model_max_budget_limiter=self._limiter,
team_id=self._token.team_id,
models=self._request_models(),
)
async def check_key_max_budget(self) -> None:
await virtual_key_max_budget_check(
valid_token=self._token, proxy_logging_obj=self._proxy_logging_obj, user_obj=self.user_object
)
def committed_spend_tracked(self) -> bool:
return committed_spend_tracked()
async def carry_committed_spend(self, read_committed: bool) -> None:
await carry_committed_budget_spend(
valid_token=self._token,
team_object=self._team_table(),
user_object=self.user_object,
end_user_object=self.end_user_object,
read_committed=partial(read_admission_committed_spend, self._spend_counter_store)
if read_committed
else None,
)
async def reserve_budget(self, skip_budget_checks: bool) -> None:
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
llm_router=llm_router,
user_api_key_auth_obj=self._token,
request=self._request,
request_data=self._request_data,
route=self._route,
llm_router=self._llm_router,
team_object=cast( # cast-ok: preserve plain team-table objects through reservation
LiteLLM_TeamTableCachedObj | None,
team_object,
self.team_object,
),
user_object=user_object,
end_user_id=end_user_id,
end_user_object=end_user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
skip_budget_checks=skip_budget_checks or sub_call,
general_settings=general_settings,
user_object=self.user_object,
end_user_id=self.end_user_id,
end_user_object=self.end_user_object,
prisma_client=self._prisma_client,
user_api_key_cache=self._user_api_key_cache,
proxy_logging_obj=self._proxy_logging_obj,
skip_budget_checks=skip_budget_checks,
general_settings=self._general_settings,
)
finally:
def release_spend_counter_batch(self) -> None:
release_spend_counter_batch()

View file

@ -0,0 +1,275 @@
from dataclasses import dataclass, replace
from typing import Final
import pytest
from litellm.proxy.auth.core.answers import EndUserScope, TokenRole
from litellm.proxy.auth.core.inputs import CheckFacts, CheckSettings, PlanSettings, RequestFacts
from litellm.proxy.auth.core.outcome import Checked, Exempt, ExemptionReason, RequestCheck
from litellm.proxy.auth.core.python_core import PythonCore
PLAN: Final = PlanSettings(
custom_auth=False,
premium_user=True,
public_routes=frozenset({"/health/liveness", "/metrics"}),
additional_public_routes=("/public/*",),
metrics_require_auth=False,
oauth2_introspection=False,
oauth2_proxy=False,
jwt_enabled=False,
oauth2_routing_overrides=(),
unauthenticated_pass_through_routes=frozenset(),
)
SETTINGS: Final = CheckSettings(
plan=PLAN,
auth_off_pass_through_routes=frozenset({"/open"}),
no_auth_dev_mode=False,
custom_auth_skips_checks=False,
has_db=True,
)
FACTS: Final = CheckFacts(
request=RequestFacts(route="/chat/completions", token="hash", is_llm_api_route=True),
force_virtual_key_checks=False,
sub_call=False,
token_has_end_user=False,
route_skips_budget_checks=False,
)
FULL_RUN: Final = (
("resolve_end_user", True),
("load_check_objects", None),
("apply_project", None),
("load_end_user_model_budget", False),
("skips_budget_checks", None),
("load_team_member_default_model_budget", False),
("merge_request_tags", None),
("bind_admission_counters", None),
("run_common_checks", False),
("check_team_model_budget", None),
("check_team_member_model_budget", None),
("check_end_user_model_budget", None),
("committed_spend_tracked", None),
("reserve_budget", False),
("release_spend_counter_batch", None),
)
class CommonCheckDenied(Exception):
pass
@dataclass
class RecordingCheckHost:
end_user: EndUserScope = EndUserScope(has_end_user=False, has_key_end_user_budget=False)
role: TokenRole = TokenRole(is_proxy_admin=False, is_builtin_admin=False)
skip_budget_checks: bool = False
committed_spend: bool = False
deny_common_checks: bool = False
calls: tuple[tuple[str, object], ...] = ()
def _record(self, name: str, payload: object = None) -> None:
self.calls = (*self.calls, (name, payload))
async def resolve_end_user(self, resolve: bool) -> EndUserScope:
self._record("resolve_end_user", resolve)
return self.end_user
async def load_check_objects(self) -> TokenRole:
self._record("load_check_objects")
return self.role
def override_admin_user(self, keep_budget_row: bool) -> None:
self._record("override_admin_user", keep_budget_row)
def apply_project(self) -> None:
self._record("apply_project")
async def apply_end_user_default_budget(self) -> None:
self._record("apply_end_user_default_budget")
async def load_end_user_model_budget(self, assign: bool) -> None:
self._record("load_end_user_model_budget", assign)
def skips_budget_checks(self) -> bool:
self._record("skips_budget_checks")
return self.skip_budget_checks
async def load_team_member_default_model_budget(self, skip_budget_checks: bool) -> None:
self._record("load_team_member_default_model_budget", skip_budget_checks)
def merge_request_tags(self) -> None:
self._record("merge_request_tags")
def bind_admission_counters(self) -> None:
self._record("bind_admission_counters")
async def run_common_checks(self, skip_budget_checks: bool) -> None:
self._record("run_common_checks", skip_budget_checks)
if self.deny_common_checks:
raise CommonCheckDenied
async def check_team_model_budget(self) -> None:
self._record("check_team_model_budget")
async def check_team_member_model_budget(self) -> None:
self._record("check_team_member_model_budget")
async def check_end_user_model_budget(self) -> None:
self._record("check_end_user_model_budget")
async def check_key_max_budget(self) -> None:
self._record("check_key_max_budget")
def committed_spend_tracked(self) -> bool:
self._record("committed_spend_tracked")
return self.committed_spend
async def carry_committed_spend(self, read_committed: bool) -> None:
self._record("carry_committed_spend", read_committed)
async def reserve_budget(self, skip_budget_checks: bool) -> None:
self._record("reserve_budget", skip_budget_checks)
def release_spend_counter_batch(self) -> None:
self._record("release_spend_counter_batch")
async def _check(
host: RecordingCheckHost, facts: CheckFacts = FACTS, settings: CheckSettings = SETTINGS
) -> RequestCheck:
return await PythonCore().authorize().check_request(host, facts, settings)
def _route(route: str, *, force: bool = False) -> CheckFacts:
return replace(FACTS, request=replace(FACTS.request, route=route), force_virtual_key_checks=force)
@pytest.mark.parametrize(
("route", "settings", "reason"),
[
("/health/liveness", SETTINGS, ExemptionReason.PUBLIC_ROUTE),
("/metrics", SETTINGS, ExemptionReason.PUBLIC_ROUTE),
("/public/docs", SETTINGS, ExemptionReason.PUBLIC_ROUTE),
("/open", SETTINGS, ExemptionReason.PASS_THROUGH_AUTH_OFF),
("/chat/completions", replace(SETTINGS, no_auth_dev_mode=True), ExemptionReason.NO_AUTH_DEV_MODE),
("/chat/completions", replace(SETTINGS, custom_auth_skips_checks=True), ExemptionReason.CUSTOM_AUTH),
],
)
async def test_exempt_requests_skip_every_check_and_do_no_io(
route: str, settings: CheckSettings, reason: ExemptionReason
) -> None:
host: Final = RecordingCheckHost()
assert await _check(host, _route(route), settings) == Exempt(reason)
assert host.calls == ()
async def test_metrics_is_checked_when_the_proxy_requires_auth_for_it() -> None:
settings: Final = replace(SETTINGS, plan=replace(PLAN, metrics_require_auth=True))
host: Final = RecordingCheckHost()
assert await _check(host, _route("/metrics"), settings) == Checked()
assert ("run_common_checks", False) in host.calls
async def test_additional_public_routes_need_a_premium_proxy() -> None:
settings: Final = replace(SETTINGS, plan=replace(PLAN, premium_user=False))
assert await _check(RecordingCheckHost(), _route("/public/docs"), settings) == Checked()
@pytest.mark.parametrize("route", ["/health/liveness", "/open"])
async def test_forced_virtual_key_checks_override_every_exemption(route: str) -> None:
settings: Final = replace(SETTINGS, no_auth_dev_mode=True, custom_auth_skips_checks=True)
host: Final = RecordingCheckHost()
assert await _check(host, _route(route, force=True), settings) == Checked()
assert host.calls == FULL_RUN
async def test_a_request_runs_every_check_in_order() -> None:
host: Final = RecordingCheckHost()
assert await _check(host) == Checked()
assert host.calls == FULL_RUN
async def test_the_end_user_is_resolved_only_when_the_token_has_none() -> None:
host: Final = RecordingCheckHost()
_ = await _check(host, replace(FACTS, token_has_end_user=True))
assert host.calls[0] == ("resolve_end_user", False)
@pytest.mark.parametrize(("is_builtin_admin", "keep_budget_row"), [(False, True), (True, False)])
async def test_a_proxy_admin_token_overrides_the_user_row_keeping_its_budget_unless_built_in(
is_builtin_admin: bool, keep_budget_row: bool
) -> None:
host: Final = RecordingCheckHost(role=TokenRole(is_proxy_admin=True, is_builtin_admin=is_builtin_admin))
_ = await _check(host)
assert host.calls[1:4] == (
("load_check_objects", None),
("override_admin_user", keep_budget_row),
("apply_project", None),
)
@pytest.mark.parametrize(
("end_user", "has_db", "applied"),
[
(EndUserScope(has_end_user=True, has_key_end_user_budget=True), True, True),
(EndUserScope(has_end_user=True, has_key_end_user_budget=True), False, False),
(EndUserScope(has_end_user=True, has_key_end_user_budget=False), True, False),
(EndUserScope(has_end_user=False, has_key_end_user_budget=True), True, False),
],
)
async def test_the_key_end_user_default_budget_applies_only_to_a_resolved_end_user_with_a_database(
end_user: EndUserScope, has_db: bool, applied: bool
) -> None:
host: Final = RecordingCheckHost(end_user=end_user)
_ = await _check(host, settings=replace(SETTINGS, has_db=has_db))
assert (("apply_end_user_default_budget", None) in host.calls) == applied
assert ("load_end_user_model_budget", end_user.has_end_user) in host.calls
async def test_budget_exempt_requests_skip_model_budgets_and_reservation() -> None:
host: Final = RecordingCheckHost(skip_budget_checks=True)
_ = await _check(host)
names: Final = tuple(name for name, _ in host.calls)
assert ("load_team_member_default_model_budget", True) in host.calls
assert ("run_common_checks", True) in host.calls
assert ("reserve_budget", True) in host.calls
assert "check_team_model_budget" not in names
async def test_routes_that_skip_budgets_keep_only_the_team_model_budget() -> None:
host: Final = RecordingCheckHost(committed_spend=True)
_ = await _check(host, replace(FACTS, route_skips_budget_checks=True))
names: Final = tuple(name for name, _ in host.calls)
assert ("load_team_member_default_model_budget", True) in host.calls
assert "check_team_model_budget" in names
assert "check_team_member_model_budget" not in names
assert "check_end_user_model_budget" not in names
assert ("carry_committed_spend", False) in host.calls
@pytest.mark.parametrize(("is_llm_api_route", "checked"), [(True, True), (False, False)])
async def test_a_sub_call_checks_the_key_budget_on_llm_routes_and_reserves_nothing(
is_llm_api_route: bool, checked: bool
) -> None:
facts: Final = replace(FACTS, sub_call=True, request=replace(FACTS.request, is_llm_api_route=is_llm_api_route))
host: Final = RecordingCheckHost()
_ = await _check(host, facts)
assert (("check_key_max_budget", None) in host.calls) == checked
assert ("reserve_budget", True) in host.calls
async def test_committed_spend_is_carried_before_reservation_when_tracked() -> None:
host: Final = RecordingCheckHost(committed_spend=True)
_ = await _check(host)
assert host.calls[-4:] == (
("committed_spend_tracked", None),
("carry_committed_spend", True),
("reserve_budget", False),
("release_spend_counter_batch", None),
)
async def test_the_spend_counter_batch_is_released_when_a_check_denies() -> None:
host: Final = RecordingCheckHost(deny_common_checks=True)
with pytest.raises(CommonCheckDenied):
_ = await _check(host)
assert host.calls[-2:] == (("run_common_checks", False), ("release_spend_counter_batch", None))