From de9287f529c9e5fde16213290154d95cc778ec3d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 10 Nov 2026 23:38:23 +0900 Subject: [PATCH] refactor(auth): prove the JWT budget check with types instead of AST tests JwtUser now carries a WithinModelBudgets or BudgetsSkipped, which only core/model_budgets.py builds, so a JWT path that skips the user model budget check no longer type-checks. A new check, check_core_proof_constructors.py, flags those proofs being built anywhere else, and in gateway-auth-core they have private fields. A user over budget on the JWT path is refused through the same denial projection as the key path. The skip decision is one pure rule, skips_model_budgets, shared by the key, JWT and common-checks paths and mirrored in Rust with the same cases. The user's model budget is now host answer data (OwnerFacts and JwtPolicy), and one attach_user_model_budgets() projection puts it on the token for the key and auto-registered paths. The last two AST tests on the JWT branch are removed. Recording-host cases in tests/unit/proxy/auth/core cover the same behavior, and both AGENTS.md files now say invariants are held by types and behavior tests, never by tests that parse source. --- .github/workflows/test-linting.yml | 3 + litellm-rust/Cargo.lock | 1 + .../crates/gateway-auth-core/AGENTS.md | 2 + .../crates/gateway-auth-core/Cargo.toml | 2 + .../crates/gateway-auth-core/src/answers.rs | 16 +- .../crates/gateway-auth-core/src/host.rs | 23 ++- .../crates/gateway-auth-core/src/identity.rs | 7 +- .../crates/gateway-auth-core/src/lib.rs | 11 +- .../gateway-auth-core/src/model_budgets.rs | 21 +++ .../crates/gateway-auth-core/src/outcome.rs | 5 +- .../gateway-auth-core/tests/model_budgets.rs | 21 +++ litellm/proxy/auth/auth_denial.py | 5 +- litellm/proxy/auth/core/AGENTS.md | 2 + litellm/proxy/auth/core/answers.py | 14 ++ litellm/proxy/auth/core/host.py | 14 +- litellm/proxy/auth/core/identity.py | 5 +- litellm/proxy/auth/core/jwt.py | 26 +++- litellm/proxy/auth/core/model_budgets.py | 22 ++- litellm/proxy/auth/core/outcome.py | 7 + litellm/proxy/auth/core/python_core.py | 28 ++-- litellm/proxy/auth/core/request_checks.py | 3 +- litellm/proxy/auth/user_api_key_auth.py | 142 +++++++++--------- .../check_core_proof_constructors.py | 43 ++++++ tests/unit/proxy/auth/core/test_jwt.py | 84 ++++++++--- .../proxy/auth/core/test_model_budgets.py | 22 ++- .../unit/proxy/auth/core/test_python_core.py | 48 ++++-- .../proxy/auth/core/test_request_checks.py | 10 +- .../unit/proxy/auth/test_user_api_key_auth.py | 78 +--------- .../test_user_api_key_auth_request_flow.py | 7 +- .../test_zero_cost_model_budget_bypass.py | 9 +- 30 files changed, 442 insertions(+), 239 deletions(-) create mode 100644 litellm-rust/crates/gateway-auth-core/src/model_budgets.rs create mode 100644 litellm-rust/crates/gateway-auth-core/tests/model_budgets.rs create mode 100644 tests/code_coverage_tests/check_core_proof_constructors.py diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index fd1d2973951..72288e7fbb5 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -611,6 +611,9 @@ jobs: - name: check_telemetry_imports run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py + - name: check_core_proof_constructors + run: uv run --no-sync python ./tests/code_coverage_tests/check_core_proof_constructors.py + - name: ban_copy_deepcopy_kwargs run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 579808c9c6a..3245e75d19d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3882,6 +3882,7 @@ version = "0.1.0" dependencies = [ "ipnet", "litellm-secret-value", + "rstest", "serde", "serde_json", "strum", diff --git a/litellm-rust/crates/gateway-auth-core/AGENTS.md b/litellm-rust/crates/gateway-auth-core/AGENTS.md index 3d2c89d2bb6..dd75e64d4b9 100644 --- a/litellm-rust/crates/gateway-auth-core/AGENTS.md +++ b/litellm-rust/crates/gateway-auth-core/AGENTS.md @@ -7,3 +7,5 @@ The crate is pure. It depends on no IO, runtime, HTTP, database or Python crate, Each Python protocol is a trait here. A core hands out its facets through accessors whose types are associated types on `AuthCore`, the same composition as the Python accessors. Host methods that do IO are `-> impl Future + Send` so facets run under Axum and `tokio::spawn` without `#[async_trait]` boxing. Values a host returns as untyped JSON in Python, such as introspection fields, are `serde_json::Value` here. A Python union of an identity and a denial is a `Result` here, as in `KeyAuthentication` `NativeCore` implements `AuthCore` with every facet `Unported`, which panics if called, the same as `RustCore` in Python. The crate exists so the types stay in step with Python. A facet leaves `Unported` only together with the Python change that routes the proxy to it, and both sides then share one test suite of cases + +Invariants are enforced by types plus behavior tests, never by tests that parse source, and they hold for the Python and Rust cores alike. A proof that a check ran, such as `WithinModelBudgets` or `BudgetsSkipped`, has a private field so only its facet's module builds it. Pure rules shared with Python, such as `skips_model_budgets`, keep the same inline cases on both sides, in `tests/` here diff --git a/litellm-rust/crates/gateway-auth-core/Cargo.toml b/litellm-rust/crates/gateway-auth-core/Cargo.toml index 4cd1ee1b636..4f336662988 100644 --- a/litellm-rust/crates/gateway-auth-core/Cargo.toml +++ b/litellm-rust/crates/gateway-auth-core/Cargo.toml @@ -13,3 +13,5 @@ serde.workspace = true serde_json.workspace = true strum.workspace = true +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/gateway-auth-core/src/answers.rs b/litellm-rust/crates/gateway-auth-core/src/answers.rs index 88932cce2ed..5d35ed5d9f2 100644 --- a/litellm-rust/crates/gateway-auth-core/src/answers.rs +++ b/litellm-rust/crates/gateway-auth-core/src/answers.rs @@ -28,9 +28,22 @@ pub struct KeyFacts { pub budget_throttle_pct: Option, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct UserModelBudgets { + pub limits: serde_json::Map, +} + +#[derive(Clone, Debug, PartialEq, Eq)] pub struct OwnerFacts { pub deactivated: bool, + pub model_budgets: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BudgetExemption { + pub route_skips_budgets: bool, + pub mcp_discovery: bool, + pub model_is_free: bool, } #[derive(Clone, Copy, Debug, PartialEq)] @@ -104,6 +117,7 @@ pub struct JwtPolicy { pub is_proxy_admin: bool, pub user_id: Option, pub owner_deactivated: bool, + pub user_model_budgets: Option, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/gateway-auth-core/src/host.rs b/litellm-rust/crates/gateway-auth-core/src/host.rs index b65a2a3214b..cc885500bd4 100644 --- a/litellm-rust/crates/gateway-auth-core/src/host.rs +++ b/litellm-rust/crates/gateway-auth-core/src/host.rs @@ -2,9 +2,9 @@ use std::{error::Error, future::Future}; use crate::{ answers::{ - AutoRegistration, EndUserScope, FallbackTarget, Introspection, JwtPolicy, KeyFacts, - KeyMapping, ModelBudgetScope, OwnerFacts, PresentedKey, TeamMember, TeamMemberBudget, - TeamMemberModelBudgetScope, TokenRole, + AutoRegistration, BudgetExemption, EndUserScope, FallbackTarget, Introspection, JwtPolicy, + KeyFacts, KeyMapping, ModelBudgetScope, OwnerFacts, PresentedKey, TeamMember, + TeamMemberBudget, TeamMemberModelBudgetScope, TokenRole, UserModelBudgets, }, identity::CustomAuthResult, inputs::RequestFacts, @@ -84,6 +84,7 @@ pub trait KeyHost: Sync { &self, user_id: &str, ) -> impl Future, Self::Error>> + Send; + fn attach_user_model_budgets(&self, budgets: Option<&UserModelBudgets>); } pub trait AuthorizeHost: Sync { @@ -93,7 +94,9 @@ pub trait AuthorizeHost: Sync { fn model_budgets(&self) -> Self::ModelBudgets; fn enforce_model_access(&self) -> impl Future> + Send; - fn check_budget_exemption(&self) -> impl Future> + Send; + fn read_budget_exemption( + &self, + ) -> impl Future> + Send; fn read_team_member_budget( &self, member: &TeamMember, @@ -120,6 +123,7 @@ pub trait AuthorizeHost: Sync { pub trait JwtHost: Sync { type Error: Error + Send + Sync + 'static; + type ModelBudgets: ModelBudgetHost; fn decode_claims(&self) -> impl Future> + Send; fn is_managed_agent(&self) -> impl Future> + Send; @@ -128,9 +132,12 @@ pub trait JwtHost: Sync { fn run_jwt_policy(&self) -> impl Future> + Send; fn remember_proxy_admin(&self) -> impl Future> + Send; fn auto_register(&self) -> impl Future> + Send; - fn skips_budget_checks(&self) -> bool; + fn attach_user_model_budgets(&self, budgets: Option<&UserModelBudgets>); + fn read_budget_exemption( + &self, + ) -> impl Future> + Send; fn load_project(&self) -> impl Future> + Send; - fn check_user_model_budgets(&self) -> impl Future> + Send; + fn model_budgets(&self) -> Self::ModelBudgets; } pub trait OAuth2IntrospectionHost: Sync { @@ -172,7 +179,9 @@ pub trait RequestCheckHost: Sync { &self, assign: bool, ) -> impl Future> + Send; - fn skips_budget_checks(&self) -> bool; + fn read_budget_exemption( + &self, + ) -> impl Future> + Send; fn load_team_member_default_model_budget( &self, skip_budget_checks: bool, diff --git a/litellm-rust/crates/gateway-auth-core/src/identity.rs b/litellm-rust/crates/gateway-auth-core/src/identity.rs index 412aab0c947..a3ac7ffe33e 100644 --- a/litellm-rust/crates/gateway-auth-core/src/identity.rs +++ b/litellm-rust/crates/gateway-auth-core/src/identity.rs @@ -4,7 +4,7 @@ use litellm_secret_value::SecretValue; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::answers::KeyFacts; +use crate::{answers::KeyFacts, model_budgets::ModelBudgetProof, outcome::ModelBudgetExceeded}; #[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] pub enum AnonymousReason { @@ -20,7 +20,7 @@ pub struct Anonymous { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum JwtIdentity { ProxyAdmin, - User, + User { budgets: ModelBudgetProof }, Key, } @@ -123,11 +123,12 @@ pub enum AuthDenial { OAuth2Proxy(OAuth2ProxyDenial), } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq)] pub enum JwtDenial { ManagedAgentNeedsDirectJwt, OwnerDeactivated { user_id: Option }, KeyMappingFailed { reason: String }, + ModelBudgetExceeded(ModelBudgetExceeded), } pub type JwtAuthentication = Result; diff --git a/litellm-rust/crates/gateway-auth-core/src/lib.rs b/litellm-rust/crates/gateway-auth-core/src/lib.rs index 5742fbb93c3..3c5ec6b2960 100644 --- a/litellm-rust/crates/gateway-auth-core/src/lib.rs +++ b/litellm-rust/crates/gateway-auth-core/src/lib.rs @@ -3,14 +3,15 @@ mod facets; mod host; mod identity; mod inputs; +mod model_budgets; mod native; mod outcome; mod validators; pub use answers::{ - AutoRegistration, EndUserScope, FallbackTarget, Introspection, JwtPolicy, KeyFacts, KeyMapping, - KeyModelBudgets, ModelBudgetScope, OwnerFacts, PresentedKey, TeamMember, TeamMemberBudget, - TeamMemberModelBudgetScope, TokenRole, + AutoRegistration, BudgetExemption, EndUserScope, FallbackTarget, Introspection, JwtPolicy, + KeyFacts, KeyMapping, KeyModelBudgets, ModelBudgetScope, OwnerFacts, PresentedKey, TeamMember, + TeamMemberBudget, TeamMemberModelBudgetScope, TokenRole, UserModelBudgets, }; pub use facets::{ AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore, @@ -30,9 +31,11 @@ pub use inputs::{ OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, RoutingOverride, }; +pub use model_budgets::{ + BudgetsSkipped, ModelBudgetProof, WithinModelBudgets, skips_model_budgets, +}; pub use native::{NativeCore, Unported}; 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/model_budgets.rs b/litellm-rust/crates/gateway-auth-core/src/model_budgets.rs new file mode 100644 index 00000000000..b57f03c4669 --- /dev/null +++ b/litellm-rust/crates/gateway-auth-core/src/model_budgets.rs @@ -0,0 +1,21 @@ +use crate::answers::BudgetExemption; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WithinModelBudgets { + _checked: (), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BudgetsSkipped { + _skipped: (), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelBudgetProof { + Within(WithinModelBudgets), + Skipped(BudgetsSkipped), +} + +pub fn skips_model_budgets(exemption: &BudgetExemption) -> bool { + exemption.route_skips_budgets || exemption.mcp_discovery || exemption.model_is_free +} diff --git a/litellm-rust/crates/gateway-auth-core/src/outcome.rs b/litellm-rust/crates/gateway-auth-core/src/outcome.rs index 9060ee8be67..ba9448c3c65 100644 --- a/litellm-rust/crates/gateway-auth-core/src/outcome.rs +++ b/litellm-rust/crates/gateway-auth-core/src/outcome.rs @@ -1,3 +1,5 @@ +use crate::model_budgets::WithinModelBudgets; + #[derive(Clone, Debug, PartialEq)] pub struct ModelBudgetExceeded { pub message: String, @@ -7,9 +9,6 @@ pub struct ModelBudgetExceeded { pub entity_id: Option, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct WithinModelBudgets; - pub type ModelBudgetResult = Result; #[derive(Clone, Debug, PartialEq)] diff --git a/litellm-rust/crates/gateway-auth-core/tests/model_budgets.rs b/litellm-rust/crates/gateway-auth-core/tests/model_budgets.rs new file mode 100644 index 00000000000..f69facb3651 --- /dev/null +++ b/litellm-rust/crates/gateway-auth-core/tests/model_budgets.rs @@ -0,0 +1,21 @@ +use litellm_gateway_auth_core::{BudgetExemption, skips_model_budgets}; +use rstest::rstest; + +#[rstest] +#[case(false, false, false, false)] +#[case(true, false, false, true)] +#[case(false, true, false, true)] +#[case(false, false, true, true)] +fn model_budgets_are_skipped_for_budget_free_routes_mcp_discovery_and_free_models( + #[case] route_skips_budgets: bool, + #[case] mcp_discovery: bool, + #[case] model_is_free: bool, + #[case] skips: bool, +) { + let exemption = BudgetExemption { + route_skips_budgets, + mcp_discovery, + model_is_free, + }; + assert_eq!(skips_model_budgets(&exemption), skips); +} diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index 6ec291db521..1eac0f7451c 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -9,7 +9,6 @@ from litellm.proxy.auth.auth_checks import key_budget_exceeded_error from litellm.proxy.auth.auth_utils import abbreviate_api_key from litellm.proxy.auth.core.identity import ( ExpiredKey, - JwtDenial, JwtKeyMappingFailed, JwtOwnerDeactivated, KeyBlocked, @@ -153,7 +152,9 @@ def denial_exception( assert_never(denial) -def jwt_denial_exception(denial: JwtDenial) -> HTTPException: +def jwt_denial_exception( + denial: ManagedAgentNeedsDirectJwt | JwtOwnerDeactivated | JwtKeyMappingFailed, +) -> HTTPException: match denial: case ManagedAgentNeedsDirectJwt(): return HTTPException(403, "Managed agents require direct JWT authentication without virtual-key mapping") diff --git a/litellm/proxy/auth/core/AGENTS.md b/litellm/proxy/auth/core/AGENTS.md index 11e1e74e11a..547c33879ac 100644 --- a/litellm/proxy/auth/core/AGENTS.md +++ b/litellm/proxy/auth/core/AGENTS.md @@ -39,3 +39,5 @@ Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exce ## Tests `tests/unit/proxy/auth/core` runs a core against a host that answers from a script and records each call. Tests assert both the result and the host calls in order, never core internals. + +Invariants are enforced by types plus behavior tests, never by tests that parse source (`ast`, `inspect.getsource`), and they must hold for the Python and the Rust host alike. A result that proves a check ran is its own type: `JwtUser` carries a `WithinModelBudgets` or `BudgetsSkipped`, which only `core/model_budgets.py` builds (`tests/code_coverage_tests/check_core_proof_constructors.py` flags any other construction), so a JWT path that skips the budget check does not type-check. Data every path must carry, like the user's model budget, is a host answer handed to one projection method that takes it as a required argument. diff --git a/litellm/proxy/auth/core/answers.py b/litellm/proxy/auth/core/answers.py index c381997dabc..6d812c09400 100644 --- a/litellm/proxy/auth/core/answers.py +++ b/litellm/proxy/auth/core/answers.py @@ -36,9 +36,22 @@ class KeyFacts: budget_throttle_pct: float | None +@dataclass(frozen=True, slots=True) +class UserModelBudgets: + limits: Mapping[str, object] + + @dataclass(frozen=True, slots=True) class OwnerFacts: deactivated: bool + model_budgets: UserModelBudgets | None + + +@dataclass(frozen=True, slots=True) +class BudgetExemption: + route_skips_budgets: bool + mcp_discovery: bool + model_is_free: bool @dataclass(frozen=True, slots=True) @@ -121,6 +134,7 @@ class JwtPolicy: is_proxy_admin: bool user_id: str | None owner_deactivated: bool + user_model_budgets: UserModelBudgets | None @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/auth/core/host.py b/litellm/proxy/auth/core/host.py index b5df9556a67..a6bfa0af801 100644 --- a/litellm/proxy/auth/core/host.py +++ b/litellm/proxy/auth/core/host.py @@ -2,6 +2,7 @@ from typing import Protocol from litellm.proxy.auth.core.answers import ( AutoRegistration, + BudgetExemption, EndUserScope, FallbackTarget, Introspection, @@ -15,6 +16,7 @@ from litellm.proxy.auth.core.answers import ( TeamMemberBudget, TeamMemberModelBudgetScope, TokenRole, + UserModelBudgets, ) from litellm.proxy.auth.core.identity import CustomAuthResult from litellm.proxy.auth.core.inputs import RequestFacts @@ -60,13 +62,15 @@ class KeyHost(Protocol): async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: ... + class AuthorizeHost(Protocol): def model_budgets(self) -> ModelBudgetHost: ... async def enforce_model_access(self) -> None: ... - async def check_budget_exemption(self) -> bool: ... + async def read_budget_exemption(self) -> BudgetExemption: ... async def read_team_member_budget(self, member: TeamMember) -> TeamMemberBudget | None: ... @@ -102,11 +106,13 @@ class JwtHost(Protocol): async def auto_register(self) -> AutoRegistration: ... - def skips_budget_checks(self) -> bool: ... + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: ... + + async def read_budget_exemption(self) -> BudgetExemption: ... async def load_project(self) -> None: ... - async def check_user_model_budgets(self) -> None: ... + def model_budgets(self) -> ModelBudgetHost: ... class OAuth2IntrospectionHost(Protocol): @@ -136,7 +142,7 @@ class RequestCheckHost(Protocol): async def load_end_user_model_budget(self, assign: bool) -> None: ... - def skips_budget_checks(self) -> bool: ... + async def read_budget_exemption(self) -> BudgetExemption: ... async def load_team_member_default_model_budget(self, skip_budget_checks: bool) -> None: ... diff --git a/litellm/proxy/auth/core/identity.py b/litellm/proxy/auth/core/identity.py index 9e91cb9b80e..6bd84cb6d8a 100644 --- a/litellm/proxy/auth/core/identity.py +++ b/litellm/proxy/auth/core/identity.py @@ -4,6 +4,7 @@ from enum import Enum from typing import TypeAlias from litellm.proxy.auth.core.answers import KeyFacts +from litellm.proxy.auth.core.outcome import ModelBudgetExceeded, ModelBudgetProof class AnonymousReason(Enum): @@ -51,7 +52,7 @@ class JwtProxyAdmin: @dataclass(frozen=True, slots=True) class JwtUser: - pass + budgets: ModelBudgetProof @dataclass(frozen=True, slots=True) @@ -212,7 +213,7 @@ class JwtKeyMappingFailed: reason: str -JwtDenial: TypeAlias = ManagedAgentNeedsDirectJwt | JwtOwnerDeactivated | JwtKeyMappingFailed +JwtDenial: TypeAlias = ManagedAgentNeedsDirectJwt | JwtOwnerDeactivated | JwtKeyMappingFailed | ModelBudgetExceeded JwtAuthentication: TypeAlias = JwtIdentity | JwtDenial diff --git a/litellm/proxy/auth/core/jwt.py b/litellm/proxy/auth/core/jwt.py index 1f16bf221c4..459394119af 100644 --- a/litellm/proxy/auth/core/jwt.py +++ b/litellm/proxy/auth/core/jwt.py @@ -20,6 +20,13 @@ from litellm.proxy.auth.core.identity import ( ManagedAgentNeedsDirectJwt, ) from litellm.proxy.auth.core.inputs import JwtSettings +from litellm.proxy.auth.core.model_budgets import ( + check_user_model_budgets, + skip_model_budgets, + skips_model_budgets, + within_model_budgets, +) +from litellm.proxy.auth.core.outcome import ModelBudgetExceeded @dataclass(frozen=True, slots=True) @@ -64,15 +71,20 @@ async def _apply_policy(host: JwtHost, settings: JwtSettings, *, pending_auto_re match registration: case RegistrationFailed(reason=reason): return JwtKeyMappingFailed(reason=reason) - case Registered(links_existing_key=True) if settings.master_key_set: - return JwtKey() - case Registered() | NotRegistered(): + case Registered(links_existing_key=links_existing_key): + host.attach_user_model_budgets(policy.user_model_budgets) + if links_existing_key and settings.master_key_set: + return JwtKey() + case NotRegistered(): pass - skip_budget_checks: Final = host.skips_budget_checks() + exemption: Final = await host.read_budget_exemption() await host.load_project() - if not skip_budget_checks: - await host.check_user_model_budgets() - return JwtUser() + if skips_model_budgets(exemption): + return JwtUser(budgets=skip_model_budgets()) + budgets: Final = within_model_budgets(await check_user_model_budgets(host.model_budgets())) + if isinstance(budgets, ModelBudgetExceeded): + return budgets + return JwtUser(budgets=budgets) async def authenticate(host: JwtHost, settings: JwtSettings) -> JwtAuthentication: diff --git a/litellm/proxy/auth/core/model_budgets.py b/litellm/proxy/auth/core/model_budgets.py index 1c9fb32f989..d9d03450712 100644 --- a/litellm/proxy/auth/core/model_budgets.py +++ b/litellm/proxy/auth/core/model_budgets.py @@ -1,13 +1,33 @@ from collections.abc import Mapping from typing import Final -from litellm.proxy.auth.core.answers import FallbackTarget, ModelBudgetScope, TeamMemberModelBudgetScope +from litellm.proxy.auth.core.answers import ( + BudgetExemption, + FallbackTarget, + ModelBudgetScope, + TeamMemberModelBudgetScope, +) from litellm.proxy.auth.core.host import ModelBudgetHost, TeamMemberModelBudgetHost from litellm.proxy.auth.core.outcome import ( + BudgetsSkipped, ModelBudgetExceeded, + ModelBudgetResult, + WithinModelBudgets, ) +def skips_model_budgets(exemption: BudgetExemption) -> bool: + return exemption.route_skips_budgets or exemption.mcp_discovery or exemption.model_is_free + + +def skip_model_budgets() -> BudgetsSkipped: + return BudgetsSkipped() + + +def within_model_budgets(exceeded: ModelBudgetExceeded | None) -> ModelBudgetResult: + return WithinModelBudgets() if exceeded is None else exceeded + + async def check_key_and_user_model_budgets(host: ModelBudgetHost) -> ModelBudgetExceeded | None: scope: Final = await host.read_model_budget_scope() if scope.key is None or not scope.models: diff --git a/litellm/proxy/auth/core/outcome.py b/litellm/proxy/auth/core/outcome.py index fae3157e36e..ed24a13b034 100644 --- a/litellm/proxy/auth/core/outcome.py +++ b/litellm/proxy/auth/core/outcome.py @@ -41,8 +41,15 @@ class WithinModelBudgets: pass +@dataclass(frozen=True, slots=True) +class BudgetsSkipped: + pass + + ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded +ModelBudgetProof: TypeAlias = WithinModelBudgets | BudgetsSkipped + Denial: TypeAlias = ( UserAuthDisabled | MasterKeyOnlyRoute | TeamMemberBudgetExceeded | KeyBudgetExceeded | ModelBudgetExceeded ) diff --git a/litellm/proxy/auth/core/python_core.py b/litellm/proxy/auth/core/python_core.py index d1850e68f2a..fe2b24b11ec 100644 --- a/litellm/proxy/auth/core/python_core.py +++ b/litellm/proxy/auth/core/python_core.py @@ -73,19 +73,19 @@ from litellm.proxy.auth.core.model_budgets import ( check_key_and_user_model_budgets, check_team_member_model_budgets, check_user_model_budgets, + skips_model_budgets, + within_model_budgets, ) from litellm.proxy.auth.core.outcome import ( Allow, Deny, KeyBudgetExceeded, MasterKeyOnlyRoute, - ModelBudgetExceeded, ModelBudgetResult, Outcome, RequestCheck, TeamMemberBudgetExceeded, UserAuthDisabled, - WithinModelBudgets, ) from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan @@ -204,17 +204,13 @@ class PythonAuthorizeCore(AuthorizeCore): class PythonModelBudgetCore(ModelBudgetCore): async def check_key_and_user(self, host: ModelBudgetHost) -> ModelBudgetResult: - return _within_model_budgets(await check_key_and_user_model_budgets(host)) + return within_model_budgets(await check_key_and_user_model_budgets(host)) async def check_user(self, host: ModelBudgetHost) -> ModelBudgetResult: - return _within_model_budgets(await check_user_model_budgets(host)) + return within_model_budgets(await check_user_model_budgets(host)) async def check_team_member(self, host: TeamMemberModelBudgetHost) -> ModelBudgetResult: - return _within_model_budgets(await check_team_member_model_budgets(host)) - - -def _within_model_budgets(exceeded: ModelBudgetExceeded | None) -> ModelBudgetResult: - return WithinModelBudgets() if exceeded is None else exceeded + return within_model_budgets(await check_team_member_model_budgets(host)) async def _presented_key(host: KeyHost, token: str, settings: KeySettings, now: datetime) -> KeyAuthentication: @@ -241,8 +237,11 @@ async def _verify_key(host: KeyHost, now: datetime, key_label: str) -> VirtualKe key: Final = await host.inspect_key() if key.blocked: return KeyBlocked() - if key.user_id is not None and await _owner_deactivated(host, key.user_id): - return OwnerDeactivated(user_id=key.user_id) + owner: Final = await host.load_key_owner(key.user_id) if key.user_id is not None else None + if owner is not None: + host.attach_user_model_budgets(owner.model_budgets) + if owner.deactivated and key.user_id is not None: + return OwnerDeactivated(user_id=key.user_id) if key.expires is not None and key.expires < now: return ExpiredKey(expires=key.expires, now=now, key_label=key_label) if not key.has_token: @@ -250,18 +249,13 @@ async def _verify_key(host: KeyHost, now: datetime, key_label: str) -> VirtualKe return VirtualKey(key=key) -async def _owner_deactivated(host: KeyHost, user_id: str) -> bool: - owner: Final = await host.load_key_owner(user_id) - return owner is not None and owner.deactivated - - async def _authorize_virtual_key( host: AuthorizeHost, key: KeyFacts, facts: RequestFacts, settings: AuthorizeSettings ) -> Outcome: if facts.route in settings.master_key_only_routes: return Deny(MasterKeyOnlyRoute()) await host.enforce_model_access() - exempt: Final = await host.check_budget_exemption() + exempt: Final = skips_model_budgets(await host.read_budget_exemption()) if not exempt and settings.has_db and key.team_member is not None: team_member_exceeded: Final = await _check_team_member_budget(host, key.team_member) if team_member_exceeded is not None: diff --git a/litellm/proxy/auth/core/request_checks.py b/litellm/proxy/auth/core/request_checks.py index eeccad43f39..e5c3c290583 100644 --- a/litellm/proxy/auth/core/request_checks.py +++ b/litellm/proxy/auth/core/request_checks.py @@ -2,6 +2,7 @@ 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.model_budgets import skips_model_budgets from litellm.proxy.auth.core.outcome import Checked, Exempt, ExemptionReason, RequestCheck from litellm.proxy.auth.core.plan import is_public_route @@ -33,7 +34,7 @@ async def check_request(host: RequestCheckHost, facts: CheckFacts, settings: Che 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() + skip_budget_checks: Final = skips_model_budgets(await host.read_budget_exemption()) await host.load_team_member_default_model_budget( skip_budget_checks=skip_budget_checks or facts.route_skips_budget_checks ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 707f165b9c6..69c5af9cd79 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -79,6 +79,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo get_team_object, get_user_object, get_user_role, + is_mcp_discovery_request, is_model_cost_zero, is_user_proxy_admin, is_valid_fallback_model, @@ -86,7 +87,6 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo key_access_group_grants_model, key_model_aliases_for_auth_check, model_max_budget_of, - request_skips_budget_checks, resolve_and_validate_end_user_id, resolve_default_end_user_budget, route_skips_budget_checks, @@ -126,6 +126,7 @@ from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, shoul from litellm.proxy.auth.core.answers import ( AdminKey, AutoRegistration, + BudgetExemption, EndUserScope, FallbackTarget, JwtPolicy, @@ -146,6 +147,7 @@ from litellm.proxy.auth.core.answers import ( TeamMemberBudget, TeamMemberModelBudgetScope, TokenRole, + UserModelBudgets, ) from litellm.proxy.auth.core.credential import ( PresentedKeys, @@ -163,11 +165,14 @@ from litellm.proxy.auth.core.identity import ( CustomAuthResult, CustomIdentity, JwtKey, + JwtKeyMappingFailed, + JwtOwnerDeactivated, JwtProxyAdmin, JwtUser, KeyAuthentication, KeyDenial, KeyIdentity, + ManagedAgentNeedsDirectJwt, MasterKey, NoCustomIdentity, NoMasterKey, @@ -2427,6 +2432,7 @@ class _JwtHost: owner_deactivated=user_object is not None and isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False, + user_model_budgets=_user_model_budgets(user_object.model_max_budget) if user_object is not None else None, ) async def remember_proxy_admin(self) -> None: @@ -2464,12 +2470,8 @@ class _JwtHost: return RegistrationFailed(reason=auto_registered.value) if auto_registered is None: return NotRegistered() - user_object: Final = result["user_object"] auto_registered.jwt_claims = self._claims auto_registered.user_email = result["user_email"] - # The auto-registered token is built from the new key's columns, which carry no user budget. Carry over - # the already-loaded user row rather than re-reading it, or the budget check has nothing to enforce. - auto_registered.user_model_max_budget = user_object.model_max_budget if user_object is not None else None self._valid_token = auto_registered self.api_key = auto_registered.token or "" return Registered( @@ -2477,17 +2479,13 @@ class _JwtHost: or pending.behavior == UnregisteredJWTClientBehavior.LINK_ONLY ) - def skips_budget_checks(self) -> bool: - model: Final = _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, - ) - return request_skips_budget_checks( - route=self._route, model=model, llm_router=self._llm_router, request_body=self._request_data - ) + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: + # An auto-registered token is built from the new key's columns, which carry no user budget, so the + # already-loaded user row's budget is carried over instead of re-reading it. + self.token().user_model_max_budget = _model_max_budget(budgets) + + async def read_budget_exemption(self) -> BudgetExemption: + return _budget_exemption(self._request_data, self._route, self._request, self._llm_router, self.token().team_id) async def load_project(self) -> None: valid_token: Final = self.token() @@ -2503,18 +2501,15 @@ class _JwtHost: valid_token.project_metadata = project.metadata valid_token.project_alias = project.project_alias - async def check_user_model_budgets(self) -> None: - await _enforce_model_budgets( - _AUTH_CORE.model_budgets().check_user, - _ModelBudgetHost( - request=self._request, - request_data=self._request_data, - valid_token=self.token(), - route=self._route, - limiter=self._limiter, - llm_router=self._llm_router, - llm_model_list=_proxy_model_list(), - ), + def model_budgets(self) -> _ModelBudgetHost: + return _ModelBudgetHost( + request=self._request, + request_data=self._request_data, + valid_token=self.token(), + route=self._route, + limiter=self._limiter, + llm_router=self._llm_router, + llm_model_list=_proxy_model_list(), ) @@ -2781,6 +2776,31 @@ def _proxy_model_list() -> list[object] | None: _MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) +def _budget_exemption( + request_data: dict[str, object], + route: str, + request: Request | None, + llm_router: litellm.Router | None, + team_id: str | None, +) -> BudgetExemption: + model: Final = _get_model_from_request_context( + request_data=request_data, route=route, request=request, llm_router=llm_router, team_id=team_id + ) + return BudgetExemption( + route_skips_budgets=route_skips_budget_checks(route=route), + mcp_discovery=is_mcp_discovery_request(route=route, request_body=request_data), + model_is_free=is_model_cost_zero(model=model, llm_router=llm_router), + ) + + +def _user_model_budgets(model_max_budget: Mapping[str, object] | None) -> UserModelBudgets | None: + return None if model_max_budget is None else UserModelBudgets(limits=model_max_budget) + + +def _model_max_budget(budgets: UserModelBudgets | None) -> dict[str, object] | None: + return None if budgets is None else dict(budgets.limits) + + def _key_settings(master_key: str | None, prisma_client: PrismaClient | None) -> KeySettings: return KeySettings(master_key=master_key, has_db=prisma_client is not None) @@ -3043,25 +3063,20 @@ class _KeyHost: self.owner = owner if owner is None: return None - key: Final = self.key - # The joint verification-token view carries the key's columns only, so the user's own per-model budget - # reaches enforcement and the post-call increment through the row fetched here. - key.user_model_max_budget = owner.model_max_budget - return OwnerFacts(deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False) + return OwnerFacts( + deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False, + model_budgets=_user_model_budgets(owner.model_max_budget), + ) - async def check_budget_exemption(self) -> bool: + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: + # The joint verification-token view carries the key's columns only, so the user's own per-model budget + # reaches enforcement and the post-call increment through the owner row. + self.key.user_model_max_budget = _model_max_budget(budgets) + + async def read_budget_exemption(self) -> BudgetExemption: from litellm.proxy.proxy_server import llm_router - model: Final = _get_model_from_request_context( - request_data=self._request_data, - route=self._route, - request=self._request, - llm_router=llm_router, - team_id=self.key.team_id, - ) - return request_skips_budget_checks( - route=self._route, model=model, llm_router=llm_router, request_body=self._request_data - ) + return _budget_exemption(self._request_data, self._route, self._request, llm_router, self.key.team_id) async def read_team_member_budget(self, member: TeamMember) -> TeamMemberBudget | None: from litellm.proxy.proxy_server import get_current_spend @@ -3304,8 +3319,18 @@ async def user_api_key_auth_builder( case JwtKey(): valid_token = jwt_host.token() api_key = jwt_host.api_key - case jwt_denial: + case ModelBudgetExceeded() as exceeded: + raise denial_exception( + exceeded, + token=jwt_host.token().token, + route=route, + jwt_auth_enabled=False, + key=jwt_host.token(), + ) + case ManagedAgentNeedsDirectJwt() | JwtOwnerDeactivated() | JwtKeyMappingFailed() as jwt_denial: raise jwt_denial_exception(jwt_denial) + case unexpected: + assert_never(unexpected) key_plan: Final = ( credential_plan if isinstance(credential_plan, KeyValidator) @@ -4168,14 +4193,8 @@ class _RequestCheckHost: dict(end_user_model_max_budget) if end_user_model_max_budget else None ) - 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, - ) + async def read_budget_exemption(self) -> BudgetExemption: + return _budget_exemption(self._request_data, self._route, self._request, self._llm_router, self._token.team_id) async def load_team_member_default_model_budget(self, skip_budget_checks: bool) -> None: await _load_team_member_default_model_budget( @@ -4395,23 +4414,6 @@ async def _reserve_budget_after_common_checks( request.state.budget_reservation = reservation # rebind-ok: read by the release middleware -def _should_skip_budget_checks( - request_data: dict[str, object], - route: str, - request: Request | None, - llm_router: litellm.Router | None, - team_id: str | None = None, -) -> bool: - model: Final = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=team_id, - ) - return request_skips_budget_checks(route=route, model=model, llm_router=llm_router, request_body=request_data) - - def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal: """Project the resolved identity into one per-request Principal, off the key object the builder already fetched, and stamp the request network context diff --git a/tests/code_coverage_tests/check_core_proof_constructors.py b/tests/code_coverage_tests/check_core_proof_constructors.py new file mode 100644 index 00000000000..d163e83394b --- /dev/null +++ b/tests/code_coverage_tests/check_core_proof_constructors.py @@ -0,0 +1,43 @@ +import ast +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +ROOT: Final = Path(__file__).resolve().parents[2] +PROOFS: Final = { + "WithinModelBudgets": "litellm/proxy/auth/core/model_budgets.py", + "BudgetsSkipped": "litellm/proxy/auth/core/model_budgets.py", +} + + +def _called_name(call: ast.Call) -> str | None: + match call.func: + case ast.Name(id=name) | ast.Attribute(attr=name): + return name + case _: + return None + + +def _violations(path: Path) -> Iterator[str]: + relative: Final = path.relative_to(ROOT).as_posix() + tree: Final = ast.parse(path.read_text(), filename=relative) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = _called_name(node) + if name in PROOFS and PROOFS[name] != relative: + yield f"{relative}:{node.lineno}: {name} is built only by {PROOFS[name]}" + + +def main() -> int: + found: Final = tuple( + violation for path in sorted((ROOT / "litellm").rglob("*.py")) for violation in _violations(path) + ) + if found: + sys.stderr.write("\n".join(found) + "\n") + return 1 if found else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unit/proxy/auth/core/test_jwt.py b/tests/unit/proxy/auth/core/test_jwt.py index e75a3930a67..f42c5f010ea 100644 --- a/tests/unit/proxy/auth/core/test_jwt.py +++ b/tests/unit/proxy/auth/core/test_jwt.py @@ -5,14 +5,18 @@ import pytest from litellm.proxy.auth.core.answers import ( AutoRegistration, + BudgetExemption, + FallbackTarget, JwtPolicy, KeyMapping, MappedKey, + ModelBudgetScope, NoKeyMapping, NotRegistered, PendingAutoRegister, Registered, RegistrationFailed, + UserModelBudgets, ) from litellm.proxy.auth.core.identity import ( JwtAuthentication, @@ -24,12 +28,23 @@ from litellm.proxy.auth.core.identity import ( ManagedAgentNeedsDirectJwt, ) from litellm.proxy.auth.core.inputs import JwtSettings +from litellm.proxy.auth.core.model_budgets import skip_model_budgets, within_model_budgets +from litellm.proxy.auth.core.outcome import ModelBudgetExceeded from litellm.proxy.auth.core.python_core import PythonCore MAPPING: Final = JwtSettings(virtual_key_mapping=True, has_db=True, master_key_set=True) NO_MAPPING: Final = replace(MAPPING, virtual_key_mapping=False) -USER_POLICY: Final = JwtPolicy(is_proxy_admin=False, user_id="u1", owner_deactivated=False) -POLICY_TAIL: Final = ("run_jwt_policy", "skips_budget_checks", "load_project", "check_user_model_budgets") +USER_BUDGETS: Final = UserModelBudgets(limits={"gpt-4o": {"budget_limit": 1.0, "time_period": "1mo"}}) +USER_POLICY: Final = JwtPolicy( + is_proxy_admin=False, user_id="u1", owner_deactivated=False, user_model_budgets=USER_BUDGETS +) +BUDGET_CHECK: Final = ("read_model_budget_scope", "check_user_model_budget") +POLICY_TAIL: Final = ("run_jwt_policy", "read_budget_exemption", "load_project", *BUDGET_CHECK) +CHECKED: Final = JwtUser(budgets=within_model_budgets(None)) +SKIPPED: Final = JwtUser(budgets=skip_model_budgets()) +EXCEEDED: Final = ModelBudgetExceeded( + message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u1" +) MAPPED_WITHOUT_EMAIL: Final = MappedKey(key_user_id="u1", key_has_email=False, claim_user_id="u1", claim_has_email=True) @@ -40,6 +55,8 @@ class RecordingJwtHost: policy: JwtPolicy = USER_POLICY registration: AutoRegistration = NotRegistered() skip_budget_checks: bool = False + user_budget_exceeded: ModelBudgetExceeded | None = None + attached: tuple[UserModelBudgets | None, ...] = () calls: tuple[str, ...] = () def _record(self, name: str) -> None: @@ -70,15 +87,38 @@ class RecordingJwtHost: self._record("auto_register") return self.registration - def skips_budget_checks(self) -> bool: - self._record("skips_budget_checks") - return self.skip_budget_checks + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: + self._record("attach_user_model_budgets") + self.attached = (*self.attached, budgets) + + async def read_budget_exemption(self) -> BudgetExemption: + self._record("read_budget_exemption") + return BudgetExemption(route_skips_budgets=False, mcp_discovery=False, model_is_free=self.skip_budget_checks) async def load_project(self) -> None: self._record("load_project") - async def check_user_model_budgets(self) -> None: - self._record("check_user_model_budgets") + def model_budgets(self) -> "RecordingJwtHost": + return self + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read_model_budget_scope") + return ModelBudgetScope(models=("gpt-4o",), rewritable_model=None, key=None, user_has_model_budgets=True) + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("check_key_model_budget") + return None + + async def authorize_budget_fallback(self, model: str) -> bool: + self._record("authorize_budget_fallback") + return False + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite_request_model") + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("check_user_model_budget") + return self.user_budget_exceeded async def _authenticate(host: RecordingJwtHost, settings: JwtSettings = MAPPING) -> JwtAuthentication: @@ -87,14 +127,20 @@ async def _authenticate(host: RecordingJwtHost, settings: JwtSettings = MAPPING) async def test_without_key_mapping_the_jwt_policy_runs_and_enforces_the_user_model_budget() -> None: host: Final = RecordingJwtHost() - assert await _authenticate(host, NO_MAPPING) == JwtUser() + assert await _authenticate(host, NO_MAPPING) == CHECKED assert host.calls == POLICY_TAIL async def test_a_zero_cost_model_skips_the_user_model_budget_on_the_jwt_path() -> None: host: Final = RecordingJwtHost(skip_budget_checks=True) - assert await _authenticate(host, NO_MAPPING) == JwtUser() - assert host.calls == ("run_jwt_policy", "skips_budget_checks", "load_project") + assert await _authenticate(host, NO_MAPPING) == SKIPPED + assert host.calls == ("run_jwt_policy", "read_budget_exemption", "load_project") + + +async def test_a_jwt_user_over_its_model_budget_is_refused() -> None: + host: Final = RecordingJwtHost(user_budget_exceeded=EXCEEDED) + assert await _authenticate(host, NO_MAPPING) == EXCEEDED + assert host.calls == POLICY_TAIL async def test_a_managed_agent_cannot_use_key_mapping() -> None: @@ -103,10 +149,11 @@ async def test_a_managed_agent_cannot_use_key_mapping() -> None: assert host.calls == ("decode_claims", "is_managed_agent") -async def test_a_mapped_key_falls_through_to_the_key_checks_without_the_jwt_policy() -> None: +async def test_a_mapped_key_falls_through_to_the_key_checks_without_the_jwt_policy_or_budget_check() -> None: host: Final = RecordingJwtHost(mapping=replace(MAPPED_WITHOUT_EMAIL, claim_has_email=False)) assert await _authenticate(host) == JwtKey() assert host.calls == ("decode_claims", "is_managed_agent", "resolve_mapping") + assert host.attached == () @pytest.mark.parametrize( @@ -129,12 +176,12 @@ async def test_a_mapped_key_backfills_its_email_only_from_claims_for_the_same_us async def test_an_unmapped_jwt_runs_the_policy_without_registering() -> None: host: Final = RecordingJwtHost() - assert await _authenticate(host) == JwtUser() + assert await _authenticate(host) == CHECKED assert host.calls == ("decode_claims", "is_managed_agent", "resolve_mapping", *POLICY_TAIL) async def test_a_deactivated_owner_is_refused_before_admin_or_registration() -> None: - policy: Final = JwtPolicy(is_proxy_admin=True, user_id="u1", owner_deactivated=True) + policy: Final = replace(USER_POLICY, is_proxy_admin=True, owner_deactivated=True) host: Final = RecordingJwtHost(mapping=PendingAutoRegister(), policy=policy) assert await _authenticate(host) == JwtOwnerDeactivated(user_id="u1") assert host.calls[-1] == "run_jwt_policy" @@ -152,7 +199,7 @@ async def test_a_proxy_admin_skips_registration_and_budgets_and_is_remembered_wh @pytest.mark.parametrize(("has_db", "registers"), [(True, True), (False, False)]) async def test_a_pending_mapping_auto_registers_only_with_a_database(has_db: bool, registers: bool) -> None: host: Final = RecordingJwtHost(mapping=PendingAutoRegister()) - assert await _authenticate(host, replace(MAPPING, has_db=has_db)) == JwtUser() + assert await _authenticate(host, replace(MAPPING, has_db=has_db)) == CHECKED assert ("auto_register" in host.calls) == registers @@ -166,9 +213,9 @@ async def test_a_failed_auto_registration_refuses_the_request() -> None: ("registration", "master_key_set", "expected"), [ (Registered(links_existing_key=True), True, JwtKey()), - (Registered(links_existing_key=True), False, JwtUser()), - (Registered(links_existing_key=False), True, JwtUser()), - (NotRegistered(), True, JwtUser()), + (Registered(links_existing_key=True), False, CHECKED), + (Registered(links_existing_key=False), True, CHECKED), + (NotRegistered(), True, CHECKED), ], ) async def test_only_a_registration_linking_an_existing_key_falls_through_to_the_key_checks( @@ -176,4 +223,5 @@ async def test_only_a_registration_linking_an_existing_key_falls_through_to_the_ ) -> None: host: Final = RecordingJwtHost(mapping=PendingAutoRegister(), registration=registration) assert await _authenticate(host, replace(MAPPING, master_key_set=master_key_set)) == expected - assert ("check_user_model_budgets" in host.calls) == isinstance(expected, JwtUser) + assert ("check_user_model_budget" in host.calls) == isinstance(expected, JwtUser) + assert host.attached == ((USER_BUDGETS,) if isinstance(registration, Registered) else ()) diff --git a/tests/unit/proxy/auth/core/test_model_budgets.py b/tests/unit/proxy/auth/core/test_model_budgets.py index d951cfc4885..7caccc6866f 100644 --- a/tests/unit/proxy/auth/core/test_model_budgets.py +++ b/tests/unit/proxy/auth/core/test_model_budgets.py @@ -5,6 +5,7 @@ from typing import Final import pytest from litellm.proxy.auth.core.answers import ( + BudgetExemption, FallbackTarget, KeyModelBudgets, ModelBudgetScope, @@ -15,6 +16,7 @@ from litellm.proxy.auth.core.outcome import ( ModelBudgetResult, WithinModelBudgets, ) +from litellm.proxy.auth.core.model_budgets import skips_model_budgets, within_model_budgets from litellm.proxy.auth.core.python_core import PythonCore Call = tuple[str, str | None] @@ -25,7 +27,7 @@ SCOPE: Final = ModelBudgetScope( key=KeyModelBudgets(fallbacks={"gpt-4o": ("gpt-4o-mini", "haiku")}), user_has_model_budgets=True, ) -ALLOWED: Final = WithinModelBudgets() +ALLOWED: Final = within_model_budgets(None) def _exceeded(entity_type: str, model: str) -> ModelBudgetExceeded: @@ -298,3 +300,21 @@ async def test_a_team_member_naming_thousands_of_models_checks_every_one(): ) assert outcome == ALLOWED assert tuple(model for name, model in calls if name == "member") == models + + +@pytest.mark.parametrize( + ("route_skips_budgets", "mcp_discovery", "model_is_free", "skips"), + [ + (False, False, False, False), + (True, False, False, True), + (False, True, False, True), + (False, False, True, True), + ], +) +def test_model_budgets_are_skipped_for_budget_free_routes_mcp_discovery_and_free_models( + route_skips_budgets: bool, mcp_discovery: bool, model_is_free: bool, skips: bool +) -> None: + exemption: Final = BudgetExemption( + route_skips_budgets=route_skips_budgets, mcp_discovery=mcp_discovery, model_is_free=model_is_free + ) + assert skips_model_budgets(exemption) is skips diff --git a/tests/unit/proxy/auth/core/test_python_core.py b/tests/unit/proxy/auth/core/test_python_core.py index ed2a4cea0bf..c5600243719 100644 --- a/tests/unit/proxy/auth/core/test_python_core.py +++ b/tests/unit/proxy/auth/core/test_python_core.py @@ -11,8 +11,10 @@ from litellm.proxy.auth.core.answers import ( ModelBudgetScope, NonAdminKey, OwnerFacts, + BudgetExemption, PresentedKey, TeamMember, + UserModelBudgets, TeamMemberBudget, ) from litellm.proxy.auth.core.identity import ( @@ -73,7 +75,7 @@ PLAIN_KEY: Final = KeyFacts( ) AUTHORIZED: Final = ( "enforce_model_access", - "check_budget_exemption", + "read_budget_exemption", "send_key_budget_alerts", "read_model_budget_scope", "load_key_context", @@ -147,9 +149,12 @@ class RecordingHandler: self._record("load_key_owner", user_id) return self.owner - async def check_budget_exemption(self) -> bool: - self._record("check_budget_exemption") - return self.exempt + async def read_budget_exemption(self) -> BudgetExemption: + self._record("read_budget_exemption") + return BudgetExemption(route_skips_budgets=False, mcp_discovery=False, model_is_free=self.exempt) + + def attach_user_model_budgets(self, budgets: UserModelBudgets | None) -> None: + self._record("attach_user_model_budgets", budgets) async def read_team_member_budget(self, member: TeamMember) -> TeamMemberBudget | None: self._record("read_team_member_budget", member) @@ -351,12 +356,17 @@ async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with(): (replace(PLAIN_KEY, blocked=True), None, KeyBlocked(), ("inspect_key",)), ( replace(PLAIN_KEY, user_id="u1"), - OwnerFacts(deactivated=True), + OwnerFacts(deactivated=True, model_budgets=None), OwnerDeactivated(user_id="u1"), - ("inspect_key", "load_key_owner"), + ("inspect_key", "load_key_owner", "attach_user_model_budgets"), ), (replace(PLAIN_KEY, has_token=False), None, KeyWithoutToken(), ("inspect_key",)), - (replace(PLAIN_KEY, user_id="u1"), OwnerFacts(deactivated=False), None, ("inspect_key", "load_key_owner")), + ( + replace(PLAIN_KEY, user_id="u1"), + OwnerFacts(deactivated=False, model_budgets=None), + None, + ("inspect_key", "load_key_owner", "attach_user_model_budgets"), + ), (replace(PLAIN_KEY, user_id="u1"), None, None, ("inspect_key", "load_key_owner")), ], ) @@ -369,6 +379,24 @@ async def test_verify_decides_whether_the_resolved_key_is_usable( assert _calls(handler) == calls +@pytest.mark.asyncio +@pytest.mark.parametrize("deactivated", [False, True]) +async def test_the_owner_model_budget_is_attached_to_the_key_even_when_the_owner_is_deactivated(deactivated: bool): + budgets: Final = UserModelBudgets(limits={"gpt-4o": {"budget_limit": 1.0, "time_period": "1mo"}}) + handler: Final = RecordingHandler( + key=replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=deactivated, model_budgets=budgets) + ) + _ = await PythonCore().keys().verify(handler, _facts("hash"), NOW) + assert handler.calls[2] == ("attach_user_model_budgets", budgets) + + +@pytest.mark.asyncio +async def test_a_key_without_an_owner_row_keeps_its_own_model_budget(): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, user_id="u1"), owner=None) + _ = await PythonCore().keys().verify(handler, _facts("hash"), NOW) + assert "attach_user_model_budgets" not in _calls(handler) + + @pytest.mark.asyncio async def test_verify_labels_an_expired_key_with_the_hash_it_was_given(): expired: Final = replace(PLAIN_KEY, expires=NOW - timedelta(seconds=1)) @@ -389,7 +417,9 @@ async def test_expired_key_over_its_team_member_budget_is_denied_as_expired_with @pytest.mark.asyncio async def test_deactivated_owner_is_denied_before_model_access_is_checked(): - handler: Final = RecordingHandler(key=replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) + handler: Final = RecordingHandler( + key=replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True, model_budgets=None) + ) authentication, calls = await _authenticate("sk-x", handler=handler) assert authentication == OwnerDeactivated(user_id="u1") assert "enforce_model_access" not in calls @@ -450,7 +480,7 @@ async def test_exempt_request_skips_every_budget_check(): ) assert await _authorize_key(key, exempt=True) == ( _allowed(), - ("enforce_model_access", "check_budget_exemption", "load_key_context"), + ("enforce_model_access", "read_budget_exemption", "load_key_context"), ) diff --git a/tests/unit/proxy/auth/core/test_request_checks.py b/tests/unit/proxy/auth/core/test_request_checks.py index 95b8c72cce4..6a4b146a258 100644 --- a/tests/unit/proxy/auth/core/test_request_checks.py +++ b/tests/unit/proxy/auth/core/test_request_checks.py @@ -3,7 +3,7 @@ from typing import Final import pytest -from litellm.proxy.auth.core.answers import EndUserScope, TokenRole +from litellm.proxy.auth.core.answers import BudgetExemption, 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 @@ -39,7 +39,7 @@ FULL_RUN: Final = ( ("load_check_objects", None), ("apply_project", None), ("load_end_user_model_budget", False), - ("skips_budget_checks", None), + ("read_budget_exemption", None), ("load_team_member_default_model_budget", False), ("merge_request_tags", None), ("bind_admission_counters", None), @@ -89,9 +89,9 @@ class RecordingCheckHost: 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 read_budget_exemption(self) -> BudgetExemption: + self._record("read_budget_exemption") + return BudgetExemption(route_skips_budgets=False, mcp_discovery=False, model_is_free=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) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 9fb95fd7d1b..76a81e17088 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1736,45 +1736,6 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) -def test_every_jwt_branch_carries_the_user_model_budget(): - """ - Each JWT branch that builds or replaces `valid_token` has to put the user's - model budget on it, or the enforcement call a few lines later has nothing to - read and silently admits the request. - - The auto-register branch is the one that regressed: it REPLACES the token - built above it with a key-scoped one whose columns carry no user budget. - """ - import ast - import inspect - import textwrap - - from litellm.proxy.auth import user_api_key_auth as auth_module - - tree = ast.parse( - textwrap.dedent(inspect.getsource(auth_module._JwtHost)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module._KeyHost)) - ) - - assignments = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Assign) - and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) - ] - targets = { - t.value.id - for node in assignments - for t in node.targets - if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) - } - assert "auto_registered" in targets, ( - f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" - ) - assert "key" in targets, "the virtual-key path must carry the user's model budget" - - @pytest.mark.asyncio async def test_user_budget_lookup_tolerates_an_unreadable_user(): """ @@ -1960,41 +1921,6 @@ def test_custom_auth_attaches_the_user_budget_even_when_it_does_not_enforce(): ) -def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): - """ - A JWT that maps to an existing virtual key resolves through the resolver - store, which builds the token from the KEY row alone and therefore carries - no user-level per-model budget. The auth core sends that branch on to the - shared virtual-key checks (pinned in tests/unit/proxy/auth/core/test_jwt.py), - where the user row is loaded and its budget copied onto the token. This - pins the copy. - """ - import ast - import inspect - import textwrap - - from litellm.proxy.auth import user_api_key_auth as auth_module - - tree = ast.parse( - textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module._KeyHost)) - ) - - # Half one: the shared block copies the user row's budget onto the token. - copies_user_row = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Assign) - and any(isinstance(t, ast.Attribute) and t.attr == "user_model_max_budget" for t in node.targets) - and any(isinstance(v, ast.Attribute) and v.attr == "model_max_budget" for v in ast.walk(node.value)) - ] - assert copies_user_row, ( - "nothing copies the user row's model_max_budget onto the token, so a mapped-key " - "JWT reaches enforcement carrying the key's columns only" - ) - - def _rejected_websocket(sent: list[Message], bearer: str) -> WebSocket: async def receive() -> Message: return {"type": "websocket.connect"} @@ -2130,12 +2056,12 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( user_api_key_cache=DualCache(), litellm_jwtauth=LiteLLM_JWTAuth(), ) - user_model_budget_check: Final = AsyncMock() + user_model_budget_check: Final = AsyncMock(return_value=None) with ( patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.auth.handle_jwt.JWTAuthManager.auth_builder", return_value=jwt_response), - patch("litellm.proxy.auth.user_api_key_auth._enforce_model_budgets", user_model_budget_check), + patch("litellm.proxy.auth.core.jwt.check_user_model_budgets", user_model_budget_check), ): result: Final = await user_api_key_auth_builder( request=mock_request, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index a38e3906594..8480f5ed803 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -13649,9 +13649,10 @@ async def test_exhausted_key_budget_spares_only_zero_spend_mcp_discovery( ("/mcp-rest/tools/call", {}, False), ], ) -def test_should_skip_budget_checks_spares_only_zero_spend_mcp_discovery( +def test_model_budgets_skip_only_zero_spend_mcp_discovery( route: str, request_data: dict[str, object], expected: bool ) -> None: - from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks + from litellm.proxy.auth.core.model_budgets import skips_model_budgets + from litellm.proxy.auth.user_api_key_auth import _budget_exemption - assert _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected + assert skips_model_budgets(_budget_exemption(request_data, route, None, None, None)) is expected diff --git a/tests/unit/proxy/test_zero_cost_model_budget_bypass.py b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py index ce1dfce1646..bc28c6e9d79 100644 --- a/tests/unit/proxy/test_zero_cost_model_budget_bypass.py +++ b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py @@ -616,7 +616,7 @@ class TestEdgeCases: class TestOverBudgetRequestThroughModelGroupAlias: """The whole path a request takes, not just the predicate. - `user_api_key_auth._should_skip_budget_checks()` derives the exemption from the requested + `user_api_key_auth._budget_exemption()` and the core's `skips_model_budgets()` derive the exemption from the requested model name and `common_checks()` enforces the budgets with it, so a break anywhere between alias resolution and enforcement shows up here. See https://github.com/BerriAI/litellm/issues/35369. @@ -649,13 +649,12 @@ class TestOverBudgetRequestThroughModelGroupAlias: async def _request(self, model: str, proxy_logging) -> bool: """Run one over-budget request for `model`, deriving the exemption the way auth does.""" - from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks + from litellm.proxy.auth.core.model_budgets import skips_model_budgets + from litellm.proxy.auth.user_api_key_auth import _budget_exemption router = self._router() request_data = {"model": model} - skip_budget_checks = _should_skip_budget_checks( - request_data=request_data, route=self.ROUTE, request=None, llm_router=router - ) + skip_budget_checks = skips_model_budgets(_budget_exemption(request_data, self.ROUTE, None, router, None)) return await common_checks( request_body=request_data, team_object=None,