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.
2026-11-13 08:52:02 +00:00 · 2026-11-10 23:38:23 +09:00 · 2026-11-10 23:38:23 +09:00 · de9287f529
commit de9287f529
parent 56e68d84c1
30 changed files with 442 additions and 239 deletions
--- 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

--- 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",
--- 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
--- 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
--- 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<f64>,
 }

-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct UserModelBudgets {
+    pub limits: serde_json::Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct OwnerFacts {
    pub deactivated: bool,
+    pub model_budgets: Option<UserModelBudgets>,
+}
+
+#[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<String>,
    pub owner_deactivated: bool,
+    pub user_model_budgets: Option<UserModelBudgets>,
 }

 #[derive(Clone, Debug, PartialEq, Eq)]
--- 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<Output = Result<Option<OwnerFacts>, 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<Output = Result<(), Self::Error>> + Send;
-    fn check_budget_exemption(&self) -> impl Future<Output = Result<bool, Self::Error>> + Send;
+    fn read_budget_exemption(
+        &self,
+    ) -> impl Future<Output = Result<BudgetExemption, Self::Error>> + 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<Output = Result<(), Self::Error>> + Send;
    fn is_managed_agent(&self) -> impl Future<Output = Result<bool, Self::Error>> + Send;
@ -128,9 +132,12 @@ pub trait JwtHost: Sync {
    fn run_jwt_policy(&self) -> impl Future<Output = Result<JwtPolicy, Self::Error>> + Send;
    fn remember_proxy_admin(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
    fn auto_register(&self) -> impl Future<Output = Result<AutoRegistration, Self::Error>> + Send;
-    fn skips_budget_checks(&self) -> bool;
+    fn attach_user_model_budgets(&self, budgets: Option<&UserModelBudgets>);
+    fn read_budget_exemption(
+        &self,
+    ) -> impl Future<Output = Result<BudgetExemption, Self::Error>> + Send;
    fn load_project(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
-    fn check_user_model_budgets(&self) -> impl Future<Output = Result<(), Self::Error>> + Send;
+    fn model_budgets(&self) -> Self::ModelBudgets;
 }

 pub trait OAuth2IntrospectionHost: Sync {
@ -172,7 +179,9 @@ pub trait RequestCheckHost: Sync {
        &self,
        assign: bool,
    ) -> impl Future<Output = Result<(), Self::Error>> + Send;
-    fn skips_budget_checks(&self) -> bool;
+    fn read_budget_exemption(
+        &self,
+    ) -> impl Future<Output = Result<BudgetExemption, Self::Error>> + Send;
    fn load_team_member_default_model_budget(
        &self,
        skip_budget_checks: bool,
--- 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<String> },
    KeyMappingFailed { reason: String },
+    ModelBudgetExceeded(ModelBudgetExceeded),
 }

 pub type JwtAuthentication = Result<JwtIdentity, JwtDenial>;
--- 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};
--- a/litellm-rust/crates/gateway-auth-core/src/model_budgets.rs
+++ 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
+}
--- 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<String>,
 }

-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub struct WithinModelBudgets;
-
 pub type ModelBudgetResult = Result<WithinModelBudgets, ModelBudgetExceeded>;

 #[derive(Clone, Debug, PartialEq)]
--- a/litellm-rust/crates/gateway-auth-core/tests/model_budgets.rs
+++ 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);
+}
--- 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")
--- 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.
--- 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)
--- 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: ...

--- 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

--- 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:
--- 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:
--- 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
 )
--- 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:
--- 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
    )
--- 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
--- a/tests/code_coverage_tests/check_core_proof_constructors.py
+++ 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())
--- 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 ())
--- 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
--- 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"),
    )


--- 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)
--- 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,
--- 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
--- 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,