diff --git a/litellm-rust/crates/gateway-auth-core/src/answers.rs b/litellm-rust/crates/gateway-auth-core/src/answers.rs index efa927f6e38..fcec57d8bb8 100644 --- a/litellm-rust/crates/gateway-auth-core/src/answers.rs +++ b/litellm-rust/crates/gateway-auth-core/src/answers.rs @@ -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, +} diff --git a/litellm-rust/crates/gateway-auth-core/src/facets.rs b/litellm-rust/crates/gateway-auth-core/src/facets.rs index c5295485004..2a054e01fda 100644 --- a/litellm-rust/crates/gateway-auth-core/src/facets.rs +++ b/litellm-rust/crates/gateway-auth-core/src/facets.rs @@ -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> + Send; + fn check_request( + &self, + host: &H, + facts: &CheckFacts, + settings: &CheckSettings, + ) -> impl Future> + Send; } pub trait AuthCore { diff --git a/litellm-rust/crates/gateway-auth-core/src/host.rs b/litellm-rust/crates/gateway-auth-core/src/host.rs index 45357bd7e58..796c413d317 100644 --- a/litellm-rust/crates/gateway-auth-core/src/host.rs +++ b/litellm-rust/crates/gateway-auth-core/src/host.rs @@ -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> + Send; } + +pub trait RequestCheckHost: Sync { + type Error: Error + Send + Sync + 'static; + + fn resolve_end_user( + &self, + resolve: bool, + ) -> impl Future> + Send; + fn load_check_objects(&self) -> impl Future> + Send; + fn override_admin_user(&self, keep_budget_row: bool); + fn apply_project(&self); + fn apply_end_user_default_budget(&self) + -> impl Future> + Send; + fn load_end_user_model_budget( + &self, + assign: bool, + ) -> impl Future> + Send; + fn skips_budget_checks(&self) -> bool; + fn load_team_member_default_model_budget( + &self, + skip_budget_checks: bool, + ) -> impl Future> + Send; + fn merge_request_tags(&self); + fn bind_admission_counters(&self); + fn run_common_checks( + &self, + skip_budget_checks: bool, + ) -> impl Future> + Send; + fn check_team_model_budget(&self) -> impl Future> + Send; + fn check_team_member_model_budget( + &self, + ) -> impl Future> + Send; + fn check_end_user_model_budget(&self) -> impl Future> + Send; + fn check_key_max_budget(&self) -> impl Future> + Send; + fn committed_spend_tracked(&self) -> bool; + fn carry_committed_spend( + &self, + read_committed: bool, + ) -> impl Future> + Send; + fn reserve_budget( + &self, + skip_budget_checks: bool, + ) -> impl Future> + Send; + fn release_spend_counter_batch(&self); +} diff --git a/litellm-rust/crates/gateway-auth-core/src/inputs.rs b/litellm-rust/crates/gateway-auth-core/src/inputs.rs index a1bfc2aa0aa..715d824460b 100644 --- a/litellm-rust/crates/gateway-auth-core/src/inputs.rs +++ b/litellm-rust/crates/gateway-auth-core/src/inputs.rs @@ -73,3 +73,21 @@ pub struct AuthorizeSettings { pub master_key_only_routes: BTreeSet, 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, + pub no_auth_dev_mode: bool, + pub custom_auth_skips_checks: bool, + pub has_db: bool, +} diff --git a/litellm-rust/crates/gateway-auth-core/src/lib.rs b/litellm-rust/crates/gateway-auth-core/src/lib.rs index 21e3b53c9d2..80ca935ce25 100644 --- a/litellm-rust/crates/gateway-auth-core/src/lib.rs +++ b/litellm-rust/crates/gateway-auth-core/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/gateway-auth-core/src/native.rs b/litellm-rust/crates/gateway-auth-core/src/native.rs index fc1fc531cf1..87d3c9aa5d5 100644 --- a/litellm-rust/crates/gateway-auth-core/src/native.rs +++ b/litellm-rust/crates/gateway-auth-core/src/native.rs @@ -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 { unimplemented!("{RUNS_IN_PYTHON}") } + + async fn check_request( + &self, + _host: &H, + _facts: &CheckFacts, + _settings: &CheckSettings, + ) -> Result { + unimplemented!("{RUNS_IN_PYTHON}") + } } diff --git a/litellm-rust/crates/gateway-auth-core/src/outcome.rs b/litellm-rust/crates/gateway-auth-core/src/outcome.rs index 9ad26ce3d32..9060ee8be67 100644 --- a/litellm-rust/crates/gateway-auth-core/src/outcome.rs +++ b/litellm-rust/crates/gateway-auth-core/src/outcome.rs @@ -34,3 +34,17 @@ pub enum Outcome { Allow { budget_throttle_pct: Option }, 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, +} diff --git a/litellm/proxy/auth/core/AGENTS.md b/litellm/proxy/auth/core/AGENTS.md index 02028a4c330..fc6409e9b9e 100644 --- a/litellm/proxy/auth/core/AGENTS.md +++ b/litellm/proxy/auth/core/AGENTS.md @@ -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`. diff --git a/litellm/proxy/auth/core/answers.py b/litellm/proxy/auth/core/answers.py index c2eeb9fec71..43357f903bf 100644 --- a/litellm/proxy/auth/core/answers.py +++ b/litellm/proxy/auth/core/answers.py @@ -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 diff --git a/litellm/proxy/auth/core/core.py b/litellm/proxy/auth/core/core.py index 3c49376733f..815bcb1dc9a 100644 --- a/litellm/proxy/auth/core/core.py +++ b/litellm/proxy/auth/core/core.py @@ -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: ... diff --git a/litellm/proxy/auth/core/host.py b/litellm/proxy/auth/core/host.py index 65d223bfc3d..4a6c394b1b2 100644 --- a/litellm/proxy/auth/core/host.py +++ b/litellm/proxy/auth/core/host.py @@ -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: ... diff --git a/litellm/proxy/auth/core/inputs.py b/litellm/proxy/auth/core/inputs.py index 2a9511023ba..ddcb0248a9e 100644 --- a/litellm/proxy/auth/core/inputs.py +++ b/litellm/proxy/auth/core/inputs.py @@ -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 diff --git a/litellm/proxy/auth/core/outcome.py b/litellm/proxy/auth/core/outcome.py index 7457f7e53c6..fae3157e36e 100644 --- a/litellm/proxy/auth/core/outcome.py +++ b/litellm/proxy/auth/core/outcome.py @@ -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 diff --git a/litellm/proxy/auth/core/python_core.py b/litellm/proxy/auth/core/python_core.py index cad7a6946ff..3cb07a81b0e 100644 --- a/litellm/proxy/auth/core/python_core.py +++ b/litellm/proxy/auth/core/python_core.py @@ -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: diff --git a/litellm/proxy/auth/core/request_checks.py b/litellm/proxy/auth/core/request_checks.py new file mode 100644 index 00000000000..eeccad43f39 --- /dev/null +++ b/litellm/proxy/auth/core/request_checks.py @@ -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() diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c619edd2c50..145842c343c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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() diff --git a/tests/unit/proxy/auth/core/test_request_checks.py b/tests/unit/proxy/auth/core/test_request_checks.py new file mode 100644 index 00000000000..95b8c72cce4 --- /dev/null +++ b/tests/unit/proxy/auth/core/test_request_checks.py @@ -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))