mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
9a79b2b2a9
commit
1edf8b268c
17 changed files with 1001 additions and 401 deletions
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
60
litellm/proxy/auth/core/request_checks.py
Normal file
60
litellm/proxy/auth/core/request_checks.py
Normal 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()
|
||||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
275
tests/unit/proxy/auth/core/test_request_checks.py
Normal file
275
tests/unit/proxy/auth/core/test_request_checks.py
Normal 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))
|
||||
Loading…
Add table
Reference in a new issue